ship customization: renaming

5 files changed, 49 insertions, 12 deletions

diff --git a/__main__.py b/__main__.py
index 36cbc28..5372651 100644
--- a/__main__.py
+++ b/__main__.py
@@ -3,7 +3,8 @@ from typing import Annotated
 
 from fastapi import FastAPI, Depends
 
-from .apirouters import agents, tasks
+from .apirouters import agents, tasks, customize_ship
+from .apirouters import tasks, customize_ship
 from .modules.database import cursor, conn
 
 
@@ -12,7 +13,7 @@ async def lifespan(app: FastAPI):
     cursor.execute(
         "CREATE TABLE IF NOT EXISTS agents(primary_key INTEGER PRIMARY KEY, callsign TEXT NOT NULL UNIQUE, token TEXT NOT NULL)")
     cursor.execute(
-        "CREATE TABLE IF NOT EXISTS ships(primary_key INTEGER PRIMARY KEY, symbol TEXT NOT NULL UNIQUE, task TEXT NOT NULL, params TEXT)")
+        "CREATE TABLE IF NOT EXISTS ships(primary_key INTEGER PRIMARY KEY, symbol TEXT NOT NULL UNIQUE, task TEXT NOT NULL, params TEXT, name TEXT)")
     conn.commit()
     agents.load_agents_from_database()
     yield
@@ -21,4 +22,5 @@ async def lifespan(app: FastAPI):
 app = FastAPI(lifespan=lifespan)
 app.include_router(agents.router)
 app.include_router(tasks.router)
+app.include_router(customize_ship.router)
 
diff --git a/apirouters/customize_ship.py b/apirouters/customize_ship.py
new file mode 100644
index 0000000..7be3279
--- /dev/null
+++ b/apirouters/customize_ship.py
@@ -0,0 +1,21 @@
+from typing import Annotated
+from fastapi import APIRouter, Depends
+from pydantic import BaseModel
+
+from .agents import auth_agent, Agent
+
+router = APIRouter()
+
+
+class RenameBody(BaseModel):
+    name: str
+
+
+@router.post("/customize_ship/{ship_symbol}/rename")
+async def rename(ship_symbol: str, rename_body: RenameBody, agent: Annotated[Agent, Depends(auth_agent)]):
+    for current_ship in agent.ships.values():
+        if current_ship.symbol == ship_symbol:
+            current_ship.rename(rename_body.name)
+            return current_ship.get_data()
+
+    return {"error": "Unknown ship symbol."}
diff --git a/apirouters/tasks.py b/apirouters/tasks.py
index 806de91..293ab52 100644
--- a/apirouters/tasks.py
+++ b/apirouters/tasks.py
@@ -1,8 +1,6 @@
 from typing import Annotated
 from fastapi import APIRouter, Depends
 
-from ..modules.database import cursor, conn
-from ..modules.ships import Ship
 from ..modules import task_type
 from .agents import auth_agent, Agent
 
@@ -12,9 +10,9 @@ router = APIRouter()
 @router.get("/tasks")
 async def get_tasks(agent: Annotated[Agent, Depends(auth_agent)]):
     ships = []
-    for ship in agent.ships.values():
-        ship.load_task()
-        ships.append({"symbol": ship.symbol, "task": ship.task})
+    for current_ship in agent.ships.values():
+        current_ship.load_task()
+        ships.append(current_ship.get_data())
 
     return {"ships": ships}
 
@@ -24,7 +22,7 @@ async def get_tasks(ship_symbol: str, agent: Annotated[Agent, Depends(auth_agent
     for current_ship in agent.ships.values():
         if current_ship.symbol == ship_symbol:
             current_ship.load_task()
-            return {"symbol": current_ship.symbol, "task": current_ship.task}
+            return current_ship.get_data()
 
     return {"error": "Unknown ship symbol."}
 
@@ -37,6 +35,7 @@ async def get_tasks(ship_symbol: str, task: str, agent: Annotated[Agent, Depends
     for current_ship in agent.ships.values():
         if current_ship.symbol == ship_symbol:
             current_ship.set_task(task)
-            return {"symbol": current_ship.symbol, "task": current_ship.task}
+            return current_ship.get_data()
 
     return {"error": "Unknown ship symbol."}
+
diff --git a/modules/ships.py b/modules/ships.py
index 8dfe41a..e35cdd2 100644
--- a/modules/ships.py
+++ b/modules/ships.py
@@ -1,3 +1,5 @@
+from typing import Dict
+
 from . import task_type
 from .database import cursor, conn
 
@@ -8,20 +10,32 @@ class Ship:
         self.callsign = "-".split(symbol)[0]
 
         self.task = None
+        self.name = None
+
+    def get_data(self) -> Dict:
+        return {"name": self.name, "symbol": self.symbol, "task": self.task}
 
     def load_task(self):
-        cursor.execute("SELECT task, params FROM ships WHERE symbol = ?", (self.symbol,))
+        cursor.execute("SELECT name, task, params FROM ships WHERE symbol = ?", (self.symbol,))
         row = cursor.fetchone()
         if row is None:
-            cursor.execute("INSERT INTO ships (symbol, task, params) VALUES (?, ?, ?)", (self.symbol, task_type.IDLE, None))
+            cursor.execute("INSERT INTO ships (symbol, task, params, name) VALUES (?, ?, ?, ?)",
+                           (self.symbol, task_type.IDLE, None, self.symbol))
             conn.commit()
 
             self.task = task_type.IDLE
+            self.task = self.symbol
         else:
-            self.task = row[0]
+            self.name = row[0]
+            self.task = row[1]
 
     def set_task(self, task):
         self.task = task
         if task != task_type.ERROR:
             cursor.execute("UPDATE ships SET task = ? WHERE symbol = ?", (task, self.symbol))
             conn.commit()
+
+    def rename(self, name):
+        self.name = name
+        cursor.execute("UPDATE ships SET name = ? WHERE symbol = ?", (name, self.symbol))
+        conn.commit()
diff --git a/requirements.txt b/requirements.txt
index 13712cc..e8c7bca 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,2 @@
 fastapi[standard]
+pydantic
\ No newline at end of file