Added doc strings, type hints, general cleaning up

Author: aefarrell

atmotube/__init__.py

@@ -1,4 +1,4 @@
-from .packets import InvalidByteData
+from .packets import InvalidByteData, AtmoTubePacket
 from .packets import StatusPacket, SPS30Packet, BME280Packet, SGPC3Packet
 from .uuids import AtmoTube_GATT_UUID
 from .gatt_notify import start_gatt_notifications, get_available_services

atmotube/gatt_notify.py

@@ -17,6 +17,14 @@
 
 
 def get_available_services(client: BleakClient) -> PacketList:
+    """
+    Get the list of supported services on the connected Atmotube device.
+
+    :param client: The BleakClient instance of the connected Atmotube device
+    :type client: BleakClient
+    :return: A list of tuples containing the UUIDs and packet classes
+    :rtype: PacketList
+    """
     characteristics = [d.characteristic_uuid
                        for c in client.services.characteristics.values()
                        for d in c.descriptors]
@@ -28,6 +36,20 @@ def get_available_services(client: BleakClient) -> PacketList:
 def gatt_notify(client: BleakClient, uuid: str | AtmoTube_GATT_UUID,
                 packet_cls: AtmoTubePacket,
                 callback: Callable[[AtmoTubePacket], None]) -> Awaitable:
+    """
+    Start GATT notifications for a specific characteristic UUID.
+
+    :param client: The BleakClient instance of the connected Atmotube device
+    :type client: BleakClient
+    :param uuid: The UUID of the characteristic to notify
+    :type uuid: str | AtmoTube_GATT_UUID
+    :param packet_cls: The packet class to instantiate from the received data
+    :type packet_cls: AtmoTubePacket
+    :param callback: The callback function to call when a packet is received
+    :type callback: Callable[[AtmoTubePacket], None]
+    :return: An awaitable object representing the notification task
+    :rtype: Awaitable
+    """
     if inspect.iscoroutinefunction(callback):
         async def packet_callback(char: BleakGATTCharacteristic,
                                   data: bytearray):
@@ -46,5 +68,15 @@ async def start_gatt_notifications(
         client: BleakClient,
         callback: Callable[[AtmoTubePacket], None],
         packet_list: PacketList = ALL_PACKETS) -> None:
+    """
+    Start GATT notifications for all specified characteristics.
+
+    :param client: The BleakClient instance of the connected Atmotube device
+    :type client: BleakClient
+    :param callback: The callback function to call when a packet is received
+    :type callback: Callable[[AtmoTubePacket], None]
+    :param packet_list: The list of UUIDs and packet classes to notify
+    :type packet_list: PacketList
+    """
     await asyncio.gather(*[gatt_notify(client, uuid, packet_cls, callback)
                            for uuid, packet_cls in packet_list])

atmotube/packets.py

@@ -1,49 +1,69 @@
+from abc import abstractmethod
 from ctypes import LittleEndianStructure, c_ubyte, c_byte, c_short, c_int
 from datetime import datetime
+from typing import TypeAlias
+
+FieldList: TypeAlias = list[tuple]
 
 
 class InvalidByteData(Exception):
     pass
 
 
+# I have no idea how to make this actually an abstract sub-class of
+# LittleEndianStructure so be mindful that this class is intended to be
+# abstract. It is only exposed to the user to use as a type hint for
+# functions that accept any Atmotube packet.
 class AtmoTubePacket(LittleEndianStructure):
-    _byte_size_ = 0  # To be defined in subclasses
+    """
+    Abstract base class for Atmotube data packets.
+    """
+    _byte_size_: int = 0  # To be defined in subclasses
 
-    def __new__(cls, data, date_time=None):
+    def __new__(cls, data: bytearray, date_time: datetime | None = None):
         if len(data) != cls._byte_size_:
             raise InvalidByteData(f"Expected {cls._byte_size_} bytes, "
                                   f"got {len(data)} bytes")
         return cls.from_buffer_copy(data)
 
-    def __init__(self, data, date_time=None):
+    def __init__(self, data: bytearray, date_time: datetime | None = None):
+
         if date_time is None:
             date_time = datetime.now()
         self.date_time = date_time
         self._process_bytes()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return str(self)
 
-    def _process_bytes(self):
+    @abstractmethod
+    def __str__(self) -> str:
+        ...
+
+    @abstractmethod
+    def _process_bytes(self) -> None:
         ...
 
 
 class StatusPacket(AtmoTubePacket):
-    _fields_ = [
-                ("_pm_sensor",          c_ubyte, 1),
-                ("_error",              c_ubyte, 1),
-                ("_bonding",            c_ubyte, 1),
-                ("_charging",           c_ubyte, 1),
-                ("_charging_timer",     c_ubyte, 1),
-                ("_bit_6",              c_ubyte, 1),
-                ("_pre_heating",        c_ubyte, 1),
-                ("_bit_8",              c_ubyte, 1),
-                ("_battery",            c_ubyte, 8),
+    """
+    Represents the status packet from an Atmotube device.
+    """
+    _fields_: FieldList = [
+                ("_pm_sensor", c_ubyte, 1),
+                ("_error", c_ubyte, 1),
+                ("_bonding", c_ubyte, 1),
+                ("_charging", c_ubyte, 1),
+                ("_charging_timer", c_ubyte, 1),
+                ("_bit_6",  c_ubyte, 1),
+                ("_pre_heating", c_ubyte, 1),
+                ("_bit_8", c_ubyte, 1),
+                ("_battery", c_ubyte, 8),
     ]
 
-    _byte_size_ = 2
+    _byte_size_: int = 2
 
-    def _process_bytes(self):
+    def _process_bytes(self) -> None:
         self.pm_sensor_status = bool(self._pm_sensor)
         self.error_flag = bool(self._error)
         self.bonding_flag = bool(self._bonding)
@@ -52,7 +72,7 @@ def _process_bytes(self):
         self.pre_heating = bool(self._pre_heating)
         self.battery_level = self._battery
 
-    def __str__(self):
+    def __str__(self) -> str:
         return (f"StatusPacket(date_time={str(self.date_time)}, "
                 f"pm_sensor_status={self.pm_sensor_status}, "
                 f"error_flag={self.error_flag}, "
@@ -64,69 +84,81 @@ def __str__(self):
 
 
 class SPS30Packet(AtmoTubePacket):
-    _fields_ = [
-        ('_pm1',   c_byte*3),
+    """
+    Represents the SPS30 particulate matter sensor data packet from an
+    Atmotube device.
+    """
+    _fields_: FieldList = [
+        ('_pm1', c_byte*3),
         ('_pm2_5', c_byte*3),
-        ('_pm10',  c_byte*3),
-        ('_pm4',   c_byte*3),
+        ('_pm10', c_byte*3),
+        ('_pm4', c_byte*3),
     ]
 
-    _pack_ = 1
-    _layout_ = "ms"
-    _byte_size_ = 12
+    _pack_: int = 1
+    _layout_: str = "ms"
+    _byte_size_: int = 12
 
-    def pm_from_bytes(self, byte_array):
+    def pm_from_bytes(self, byte_array: bytearray) -> float | None:
         res = int.from_bytes(byte_array, byteorder='little', signed=True)
         return res/100.0 if res > 0 else None
 
-    def _process_bytes(self):
+    def _process_bytes(self) -> None:
         self.pm1 = self.pm_from_bytes(self._pm1)
         self.pm2_5 = self.pm_from_bytes(self._pm2_5)
         self.pm10 = self.pm_from_bytes(self._pm10)
         self.pm4 = self.pm_from_bytes(self._pm4)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return (f"SPS30Packet(date_time={str(self.date_time)}, "
                 f"pm1={self.pm1}µg/m³, pm2_5={self.pm2_5}µg/m³, "
                 f"pm10={self.pm10}µg/m³, pm4={self.pm4}µg/m³)")
 
 
 class BME280Packet(AtmoTubePacket):
-    _fields_ = [
-        ('_rh',    c_byte),
-        ('_T',     c_byte),
-        ('_P',     c_int),
+    """
+    Represents the BME280 environmental sensor data packet from an
+    Atmotube device.
+    """
+    _fields_: FieldList = [
+        ('_rh', c_byte),
+        ('_T', c_byte),
+        ('_P', c_int),
         ('_T_dec', c_short),
         ]
 
-    _pack_ = 1
-    _layout_ = "ms"
-    _byte_size_ = 8
+    _pack_: int = 1
+    _layout_: str = "ms"
+    _byte_size_: int = 8
 
-    def _process_bytes(self):
+    def _process_bytes(self) -> None:
         self.humidity = self._rh if self._rh > 0 else None
         self.temperature = self._T_dec / 100.0
         self.pressure = self._P / 100.0 if self._P > 0 else None
 
-    def __str__(self):
+    def __str__(self) -> str:
         return (f"BME280Packet(date_time={str(self.date_time)}, "
                 f"humidity={self.humidity}%, "
                 f"temperature={self.temperature}°C, "
                 f"pressure={self.pressure}mbar)")
 
 
 class SGPC3Packet(AtmoTubePacket):
-    _fields_ = [
+    """
+    Represents the SGPC3 air quality sensor data packet from an
+    Atmotube device.
+    """
+    _fields_: FieldList = [
         ('_tvoc', c_short)
     ]
 
-    _pack_ = 1
-    _layout_ = "ms"
-    _byte_size_ = 4
+    _pack_: int = 1
+    _layout_: str = "ms"
+    _byte_size_: int = 4
 
-    def _process_bytes(self):
+    def _process_bytes(self) -> None:
         self.tvoc = self._tvoc/1000.0 if self._tvoc > 0 else None
 
-    def __str__(self):
+    def __str__(self) -> str:
         return (f"SGPC3Packet(date_time={str(self.date_time)}, "
                 f"tvoc={self.tvoc}ppb)")

examples/data_logging_example.py

@@ -7,15 +7,25 @@
 
 from bleak import BleakClient, BleakScanner
 from atmotube import SPS30Packet, StatusPacket, BME280Packet, SGPC3Packet
-from atmotube import start_gatt_notifications, get_available_services
+from atmotube import AtmoTubePacket, start_gatt_notifications
+from atmotube import get_available_services
 import asyncio
 import logging
 
-ATMOTUBE = "C2:2B:42:15:30:89"  # the mac address of my Atmotube
 
+async def collect_data(mac: str, queue: asyncio.Queue,
+                       collection_time: int) -> None:
+    """
+    Connects to the Atmotube device and collects data for a specified time.
 
-async def collect_data(mac, queue, collection_time):
-    async def callback_queue(packet):
+    :param mac: The MAC address of the Atmotube device
+    :type mac: str
+    :param queue: An asyncio Queue to put the received packets into
+    :type queue: asyncio.Queue
+    :param collection_time: The duration in seconds to collect data
+    :type collection_time: int
+    """
+    async def callback_queue(packet: AtmoTubePacket) -> None:
         await queue.put(packet)
 
     device = await BleakScanner.find_device_by_address(mac)
@@ -32,7 +42,13 @@ async def callback_queue(packet):
         await queue.put(None)
 
 
-def log_packet(packet):
+def log_packet(packet: AtmoTubePacket) -> None:
+    """
+    Logs the received packet to the console.
+
+    :param packet: The packet to log
+    :type packet: AtmoTubePacket
+    """
     match packet:
         case StatusPacket():
             logging.info(f"{str(packet.date_time)} - Status Packet - "
@@ -58,12 +74,15 @@ def log_packet(packet):
             logging.info("Unknown packet type")
 
 
-def main():
-    mac = ATMOTUBE
+def main() -> None:
+    mac = "C2:2B:42:15:30:89"  # the mac address of my Atmotube
     collection_time = 60  # seconds
     queue = asyncio.Queue()
 
-    async def runner():
+    async def runner() -> None:
+        """
+        Main runner function to collect and log data.
+        """
         collector = asyncio.create_task(
             collect_data(mac, queue, collection_time)
             )