dotbot/controller: implement dotbot filtering using a query model

doc/conf.py

@@ -42,6 +42,9 @@
     ("py:class", r"dotbot_utils.*"),
     ("py:class", r"ASGIApp"),
     ("py:class", r"DispatchFunction"),
+    ("py:class", r"dotbot.models.Annotated"),
+    ("py:class", r"Query"),
+    ("py:class", r"PydanticUndefined"),
 ]
 
 # -- Options for HTML output -------------------------------------------------

dotbot/controller.py

@@ -736,19 +736,60 @@ def get_dotbots(self, query: DotBotQueryModel) -> List[DotBotModel]:
         """Returns the list of dotbots matching the query."""
         dotbots: List[DotBotModel] = []
         for dotbot in self.dotbots.values():
+            if query.address is not None and dotbot.address != query.address:
+                continue
             if (
                 query.application is not None
-                and dotbot.application != query.application
+                and dotbot.application.value != query.application
             ):
                 continue
-            if query.mode is not None and dotbot.mode != query.mode:
+            if query.status is not None and dotbot.status.value != query.status:
+                continue
+            if query.max_battery is not None and dotbot.battery is not None:
+                if dotbot.battery > query.max_battery:
+                    continue
+            if query.min_battery is not None and dotbot.battery is not None:
+                if dotbot.battery < query.min_battery:
+                    continue
+            if (
+                any(
+                    [
+                        query.max_position_x is not None,
+                        query.min_position_x is not None,
+                        query.max_position_y is not None,
+                        query.min_position_y is not None,
+                    ]
+                )
+                and dotbot.lh2_position is None
+            ):
                 continue
-            if query.status is not None and dotbot.status != query.status:
+            if dotbot.lh2_position is None and query.max_positions is not None:
                 continue
+            if dotbot.lh2_position is not None:
+                if query.max_position_x is not None:
+                    if query.max_position_x < dotbot.lh2_position.x:
+                        continue
+                if query.min_position_x is not None:
+                    if query.min_position_x > dotbot.lh2_position.x:
+                        continue
+                if query.max_position_y is not None:
+                    if query.max_position_y < dotbot.lh2_position.y:
+                        continue
+                if query.min_position_y is not None:
+                    if query.min_position_y > dotbot.lh2_position.y:
+                        continue
             _dotbot = DotBotModel(**dotbot.model_dump())
-            _dotbot.position_history = _dotbot.position_history[: query.max_positions]
+            max_positions = (
+                MAX_POSITION_HISTORY_SIZE
+                if query.max_positions is None
+                else query.max_positions
+            )
+            _dotbot.position_history = _dotbot.position_history[:max_positions]
             dotbots.append(_dotbot)
-        return sorted(dotbots, key=lambda dotbot: dotbot.address)
+        dotbots = sorted(dotbots, key=lambda dotbot: dotbot.address)
+        if query.limit is not None:
+            dotbots = dotbots[: query.limit]
+        return dotbots
 
     async def web(self):
         """Starts the web server application."""

dotbot/examples/charging_station.py

@@ -13,7 +13,9 @@
     DotBotLH2Position,
     DotBotModel,
     DotBotMoveRawCommandModel,
+    DotBotQueryModel,
     DotBotRgbLedCommandModel,
+    DotBotStatus,
     DotBotWaypoints,
     WSRgbLed,
     WSWaypoints,
@@ -52,12 +54,18 @@ async def queue_robots(
     await send_to_goal(client, ws, goals, params)
 
 
+async def fetch_active_dotbots(client: RestClient) -> List[DotBotModel]:
+    return await client.fetch_dotbots(
+        query=DotBotQueryModel(status=DotBotStatus.ACTIVE)
+    )
+
+
 async def charge_robots(
     client: RestClient,
     ws: DotBotWsClient,
     params: OrcaParams,
 ) -> None:
-    dotbots = await client.fetch_active_dotbots()
+    dotbots = await fetch_active_dotbots(client)
     remaining = order_bots(dotbots, QUEUE_HEAD_X, QUEUE_HEAD_Y)
     total_count = len(dotbots)
     # The head of the remaining should park
@@ -66,7 +74,7 @@ async def charge_robots(
     parked_count = total_count - len(remaining)
 
     while remaining or park_dotbot is not None:
-        dotbots = await client.fetch_active_dotbots()
+        dotbots = await fetch_active_dotbots(client)
 
         dotbots = [b for b in dotbots if b.address in {r.address for r in remaining}]
         remaining = order_bots(dotbots, QUEUE_HEAD_X, QUEUE_HEAD_Y)
@@ -133,7 +141,7 @@ async def send_to_goal(
     params: OrcaParams,
 ) -> None:
     while True:
-        dotbots = await client.fetch_active_dotbots()
+        dotbots = await fetch_active_dotbots(client)
         agents: List[Agent] = []
 
         for bot in dotbots:
@@ -309,7 +317,7 @@ async def main() -> None:
     port = os.getenv("DOTBOT_CONTROLLER_PORT", "8000")
     use_https = os.getenv("DOTBOT_CONTROLLER_USE_HTTPS", False)
     async with rest_client(url, port, use_https) as client:
-        dotbots = await client.fetch_active_dotbots()
+        dotbots = await fetch_active_dotbots(client)
 
         ws = DotBotWsClient(url, port)
         await ws.connect()

dotbot/joystick.py

@@ -20,7 +20,7 @@
     pydotbot_version,
 )
 from dotbot.logger import LOGGER, setup_logging
-from dotbot.models import DotBotMoveRawCommandModel
+from dotbot.models import DotBotMoveRawCommandModel, DotBotQueryModel, DotBotStatus
 from dotbot.protocol import ApplicationType
 from dotbot.rest import rest_client
 
@@ -99,12 +99,18 @@ def pos_from_joystick(self):
 
     async def fetch_active_dotbots(self):
         while 1:
-            self.dotbots = await self.api.fetch_active_dotbots()
+            self.dotbots = await self.api.fetch_active_dotbots(
+                query=DotBotQueryModel(status=DotBotStatus.ACTIVE)
+            )
             await asyncio.sleep(1)
 
     async def start(self):
         """Starts to read continuously joystick positions."""
-        asyncio.create_task(self.fetch_active_dotbots())
+        asyncio.create_task(
+            self.fetch_active_dotbots(
+                query=DotBotQueryModel(status=DotBotStatus.ACTIVE)
+            )
+        )
         while True:
             # fetch positions from joystick
             positions = self.pos_from_joystick()

dotbot/keyboard.py

@@ -12,6 +12,7 @@
 
 import click
 
+from dotbot.models import DotBotQueryModel, DotBotStatus
 from dotbot.rest import rest_client
 
 try:
@@ -236,12 +237,18 @@ async def refresh_speeds(self):
 
     async def fetch_active_dotbots(self):
         while 1:
-            self.dotbots = await self.api.fetch_active_dotbots()
+            self.dotbots = await self.api.fetch_active_dotbots(
+                query=DotBotQueryModel(status=DotBotStatus.ACTIVE)
+            )
             await asyncio.sleep(1)
 
     async def start(self):
         """Starts to continuously listen on keyboard key press/release events."""
-        asyncio.create_task(self.fetch_active_dotbots())
+        asyncio.create_task(
+            self.fetch_active_dotbots(
+                query=DotBotQueryModel(status=DotBotStatus.ACTIVE)
+            )
+        )
         asyncio.create_task(self.update_active_keys())
         while 1:
             await self.refresh_speeds()

dotbot/models.py

@@ -59,7 +59,7 @@ class DotBotLH2Position(BaseModel):
 
     x: float
     y: float
-    z: float
+    z: float = 0.0
 
 
 class DotBotControlModeModel(BaseModel):
@@ -100,11 +100,17 @@ class DotBotStatus(IntEnum):
 class DotBotQueryModel(BaseModel):
     """Model class used to filter DotBots."""
 
-    max_positions: int = MAX_POSITION_HISTORY_SIZE
+    limit: Optional[int] = None
+    address: Optional[str] = None
     application: Optional[ApplicationType] = None
-    mode: Optional[ControlModeType] = None
     status: Optional[DotBotStatus] = None
-    swarm: Optional[str] = None
+    max_battery: Optional[float] = None
+    min_battery: Optional[float] = None
+    max_positions: int = None
+    max_position_x: Optional[float] = None
+    min_position_x: Optional[float] = None
+    max_position_y: Optional[float] = None
+    min_position_y: Optional[float] = None
 
 
 class DotBotNotificationCommand(IntEnum):
@@ -179,7 +185,7 @@ class DotBotModel(BaseModel):
     waypoints_threshold: int = 100  # in mm
     position_history: List[Union[DotBotLH2Position, DotBotGPSPosition]] = []
     calibrated: int = 0x00  # Bitmask: first lighthouse = 0x01, second lighthouse = 0x02
-    battery: float = 0.0  # Voltage in Volts
+    battery: float = 3.0  # Voltage in Volts
 
 
 class WSBase(BaseModel):

dotbot/rest.py

@@ -5,13 +5,14 @@
 
 """Module containing client code to interact with the controller REST API."""
 
+import urllib.parse
 from contextlib import asynccontextmanager
-from typing import List
+from typing import List, Optional
 
 import httpx
 
 from dotbot.logger import LOGGER, setup_logging
-from dotbot.models import DotBotModel, DotBotStatus
+from dotbot.models import DotBotModel, DotBotQueryModel
 from dotbot.protocol import ApplicationType
 
 
@@ -43,11 +44,16 @@ def base_url(self):
     async def close(self):
         await self._client.aclose()
 
-    async def fetch_active_dotbots(self) -> List[DotBotModel]:
-        """Fetch active DotBots."""
+    async def fetch_dotbots(
+        self, query: Optional[DotBotQueryModel] = None
+    ) -> List[DotBotModel]:
+        """Fetch DotBots matching the query."""
         try:
+            url = f"{self.base_url}/dotbots"
+            if query is not None:
+                url += f"?{urllib.parse.urlencode(query.model_dump(exclude_none=True))}"
             response = await self._client.get(
-                f"{self.base_url}/dotbots",
+                url,
                 headers={
                     "Accept": "application/json",
                 },
@@ -60,11 +66,7 @@ async def fetch_active_dotbots(self) -> List[DotBotModel]:
                     f"Failed to fetch dotbots: {response} {response.text}"
                 )
             else:
-                return [
-                    DotBotModel(**dotbot)
-                    for dotbot in response.json()
-                    if dotbot["status"] == DotBotStatus.ACTIVE.value
-                ]
+                return [DotBotModel(**dotbot) for dotbot in response.json()]
         return []
 
     async def _send_command(self, address, application, resource, command):

dotbot/server.py

@@ -6,10 +6,16 @@
 """Module for the web server application."""
 
 import os
-from typing import List
+from typing import Annotated, List
 
 import httpx
-from fastapi import Depends, FastAPI, HTTPException, WebSocket, WebSocketDisconnect
+from fastapi import (
+    FastAPI,
+    HTTPException,
+    Query,
+    WebSocket,
+    WebSocketDisconnect,
+)
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import Response
 from fastapi.staticfiles import StaticFiles
@@ -19,6 +25,7 @@
 from dotbot import pydotbot_version
 from dotbot.logger import LOGGER
 from dotbot.models import (
+    MAX_POSITION_HISTORY_SIZE,
     DotBotMapSizeModel,
     DotBotModel,
     DotBotMoveRawCommandModel,
@@ -228,12 +235,12 @@ async def dotbot_positions_history_clear(address: str):
     summary="Return information about a dotbot given its address",
     tags=["dotbots"],
 )
-async def dotbot(address: str, query: DotBotQueryModel = Depends()):
+async def dotbot(address: str, max_positions: int = MAX_POSITION_HISTORY_SIZE):
     """Dotbot HTTP GET handler."""
     if address not in api.controller.dotbots:
         raise HTTPException(status_code=404, detail="No matching dotbot found")
     _dotbot = DotBotModel(**api.controller.dotbots[address].model_dump())
-    _dotbot.position_history = _dotbot.position_history[: query.max_positions]
+    _dotbot.position_history = _dotbot.position_history[:max_positions]
     return _dotbot
 
 
@@ -244,7 +251,7 @@ async def dotbot(address: str, query: DotBotQueryModel = Depends()):
     summary="Return the list of available dotbots",
     tags=["dotbots"],
 )
-async def dotbots(query: DotBotQueryModel = Depends()):
+async def dotbots(query: Annotated[DotBotQueryModel, Query()]):
     """Dotbots HTTP GET handler."""
     return api.controller.get_dotbots(query)