Add maximum number of items to fetch

ttu · Jan 24, 2025 · 94ac7b3 · 94ac7b3
1 parent f34deb5
commit 94ac7b3
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 11 deletions.
diff --git a/ruuvitag_sensor/adapters/bleak_ble.py b/ruuvitag_sensor/adapters/bleak_ble.py
@@ -126,13 +126,19 @@ async def get_first_data(mac: str, bt_device: str = "") -> RawData:
 
         return data or ""
 
-    async def get_history_data(self, mac: str, start_time: Optional[datetime] = None) -> List[dict]:
+    async def get_history_data(
+        self, 
+        mac: str, 
+        start_time: Optional[datetime] = None,
+        max_items: Optional[int] = None
+    ) -> List[dict]:
         """
         Get history data from a RuuviTag using GATT connection.
 
         Args:
             mac (str): MAC address of the RuuviTag
             start_time (datetime, optional): Start time for history data
+            max_items (int, optional): Maximum number of history entries to fetch
 
         Returns:
             List[dict]: List of historical sensor readings
@@ -184,13 +190,17 @@ def notification_handler(_, data: bytearray):
             log.debug("Requested history data from device %s", mac)
 
             # Wait for initial notification
-            await asyncio.wait_for(notification_received.wait(), timeout=5.0)
+            await asyncio.wait_for(notification_received.wait(), timeout=10.0)
 
             # Wait for more data
             try:
                 while True:
                     notification_received.clear()
-                    await asyncio.wait_for(notification_received.wait(), timeout=1.0)
+                    await asyncio.wait_for(notification_received.wait(), timeout=5.0)
+                    # Check if we've reached the maximum number of items
+                    if max_items and len(history_data) >= max_items:
+                        log.debug("Reached maximum number of items (%d)", max_items)
+                        break
             except asyncio.TimeoutError:
                 # No more data received for 1 second - assume transfer complete
                 pass

diff --git a/ruuvitag_sensor/ruuvi.py b/ruuvitag_sensor/ruuvi.py
@@ -346,29 +346,40 @@ def _parse_data(
         return (mac_to_send, decoded)
 
     @staticmethod
-    async def get_history_async(mac: str, start_time: Optional[datetime] = None) -> List[SensorData]:
+    async def get_history_async(
+        mac: str, 
+        start_time: Optional[datetime] = None,
+        max_items: Optional[int] = None
+    ) -> List[SensorData]:
         """
         Get history data from a RuuviTag that supports it (firmware 3.30.0+)
 
         Args:
             mac (str): MAC address of the RuuviTag
             start_time (datetime, optional): Start time for history data. If None, gets all available data
+            max_items (int, optional): Maximum number of history entries to fetch. If None, gets all available data
 
         Returns:
             List[SensorData]: List of historical sensor readings
         """
         throw_if_not_async_adapter(ble)
-        return await ble.get_history_data(mac, start_time)
+        return await ble.get_history_data(mac, start_time, max_items)
 
     @staticmethod
-    async def download_history(mac: str, start_time: Optional[datetime] = None, timeout: int = 300) -> List[SensorData]:
+    async def download_history(
+        mac: str, 
+        start_time: Optional[datetime] = None, 
+        timeout: int = 300,
+        max_items: Optional[int] = None
+    ) -> List[SensorData]:
         """
         Download history data from a RuuviTag. Requires firmware version 3.30.0 or newer.
 
         Args:
             mac (str): MAC address of the RuuviTag. On macOS use UUID instead.
             start_time (Optional[datetime]): If provided, only get data from this time onwards
-            timeout (int): Maximum time in seconds to wait for history download (default: 30)
+            timeout (int): Maximum time in seconds to wait for history download (default: 300)
+            max_items (Optional[int]): Maximum number of history entries to fetch. If None, gets all available data
 
         Returns:
             List[SensorData]: List of historical measurements, ordered by timestamp
@@ -379,11 +390,11 @@ async def download_history(mac: str, start_time: Optional[datetime] = None, time
         """
         throw_if_not_async_adapter(ble)
 
-        # if not re.match("[0-9A-F]{2}(:[0-9A-F]{2}){5}$", mac.upper()):
-        #     raise ValueError(f"Invalid MAC address: {mac}")
-
         try:
-            history = await asyncio.wait_for(ble.get_history_data(mac, start_time), timeout=timeout)
+            history = await asyncio.wait_for(
+                ble.get_history_data(mac, start_time, max_items), 
+                timeout=timeout
+            )
 
             # Sort by timestamp if present
             if history and "timestamp" in history[0]: