bring async_.Routine back and support passing an optional task group

Author: rnovatorov

examples/mqtt/pub_sub.py

@@ -6,10 +6,10 @@
 
 async def main() -> None:
     config = enapter.mqtt.Config(host="127.0.0.1", port=1883)
-    async with asyncio.TaskGroup() as tg:
-        client = enapter.mqtt.Client(tg, config=config)
-        tg.create_task(subscriber(client))
-        tg.create_task(publisher(client))
+    async with enapter.mqtt.Client(config=config) as client:
+        async with asyncio.TaskGroup() as tg:
+            tg.create_task(subscriber(client))
+            tg.create_task(publisher(client))
 
 
 async def subscriber(client: enapter.mqtt.Client) -> None:

examples/standalone/zigbee2mqtt/script.py

@@ -36,13 +36,15 @@ async def run(self):
             tg.create_task(self.properties_sender())
 
     async def consumer(self, tg):
-        client = enapter.mqtt.Client(tg, self.mqtt_client_config)
-        async with client.subscribe(self.mqtt_topic) as messages:
-            async for msg in messages:
-                try:
-                    self.telemetry = json.loads(msg.payload)
-                except json.JSONDecodeError as e:
-                    await self.log.error(f"failed to decode json payload: {e}")
+        async with enapter.mqtt.Client(
+            self.mqtt_client_config, task_group=tg
+        ) as client:
+            async with client.subscribe(self.mqtt_topic) as messages:
+                async for msg in messages:
+                    try:
+                        self.telemetry = json.loads(msg.payload)
+                    except json.JSONDecodeError as e:
+                        await self.log.error(f"failed to decode json payload: {e}")
 
     async def telemetry_sender(self):
         while True:

src/enapter/async_/__init__.py

@@ -1,3 +1,4 @@
 from .generator import generator
+from .routine import Routine
 
-__all__ = ["generator"]
+__all__ = ["generator", "Routine"]

src/enapter/async_/routine.py

@@ -0,0 +1,37 @@
+import abc
+import asyncio
+import contextlib
+from typing import Self
+
+
+class Routine(abc.ABC):
+
+    def __init__(self, task_group: asyncio.TaskGroup | None) -> None:
+        self._task_group = task_group
+        self._task: asyncio.Task | None = None
+
+    @abc.abstractmethod
+    async def _run(self) -> None:
+        pass
+
+    async def __aenter__(self) -> Self:
+        await self.start()
+        return self
+
+    async def __aexit__(self, *_) -> None:
+        await self.stop()
+
+    async def start(self) -> None:
+        if self._task is not None:
+            raise RuntimeError("already started")
+        if self._task_group is None:
+            self._task = asyncio.create_task(self._run())
+        else:
+            self._task = self._task_group.create_task(self._run())
+
+    async def stop(self) -> None:
+        if self._task is None:
+            raise RuntimeError("not started yet")
+        self._task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._task

src/enapter/mqtt/client.py

@@ -15,19 +15,18 @@
 LOGGER = logging.getLogger(__name__)
 
 
-class Client:
+class Client(async_.Routine):
 
-    def __init__(self, task_group: asyncio.TaskGroup, config: Config) -> None:
+    def __init__(
+        self, config: Config, task_group: asyncio.TaskGroup | None = None
+    ) -> None:
+        super().__init__(task_group=task_group)
         self._logger = self._new_logger(config)
         self._config = config
         self._mdns_resolver = mdns.Resolver()
         self._tls_context = self._new_tls_context(config)
         self._publisher: aiomqtt.Client | None = None
         self._publisher_connected = asyncio.Event()
-        self._task = task_group.create_task(self._run())
-
-    def cancel(self) -> None:
-        self._task.cancel()
 
     @staticmethod
     def _new_logger(config: Config) -> logging.LoggerAdapter:

src/enapter/standalone/__init__.py

@@ -1,4 +1,3 @@
-from .app import App, run
 from .config import Config
 from .device import Device
 from .device_protocol import (
@@ -10,9 +9,9 @@
     Telemetry,
 )
 from .logger import Logger
+from .run import run
 
 __all__ = [
-    "App",
     "CommandArgs",
     "CommandResult",
     "Config",

src/enapter/standalone/app.py

@@ -3,22 +3,22 @@
 import time
 import traceback
 
-from enapter import mqtt
+from enapter import async_, mqtt
 
 from .device_protocol import DeviceProtocol
 
 
-class DeviceDriver:
+class DeviceDriver(async_.Routine):
 
     def __init__(
         self,
-        task_group: asyncio.TaskGroup,
         device_channel: mqtt.api.DeviceChannel,
         device: DeviceProtocol,
+        task_group: asyncio.TaskGroup | None,
     ) -> None:
+        super().__init__(task_group=task_group)
         self._device_channel = device_channel
         self._device = device
-        self._task = task_group.create_task(self._run())
 
     async def _run(self) -> None:
         async with asyncio.TaskGroup() as tg:

src/enapter/standalone/device_driver.py

@@ -0,0 +1,43 @@
+import asyncio
+import contextlib
+
+from enapter import log, mqtt
+
+from .config import Config
+from .device_driver import DeviceDriver
+from .device_protocol import DeviceProtocol
+from .ucm import UCM
+
+
+async def run(device: DeviceProtocol) -> None:
+    log.configure(level=log.LEVEL or "info")
+    config = Config.from_env()
+    async with contextlib.AsyncExitStack() as stack:
+        task_group = await stack.enter_async_context(asyncio.TaskGroup())
+        mqtt_client = await stack.enter_async_context(
+            mqtt.Client(config=config.communication.mqtt, task_group=task_group)
+        )
+        _ = await stack.enter_async_context(
+            DeviceDriver(
+                device_channel=mqtt.api.DeviceChannel(
+                    client=mqtt_client,
+                    hardware_id=config.communication.hardware_id,
+                    channel_id=config.communication.channel_id,
+                ),
+                device=device,
+                task_group=task_group,
+            )
+        )
+        if config.communication.ucm_needed:
+            _ = await stack.enter_async_context(
+                DeviceDriver(
+                    device_channel=mqtt.api.DeviceChannel(
+                        client=mqtt_client,
+                        hardware_id=config.communication.hardware_id,
+                        channel_id="ucm",
+                    ),
+                    device=UCM(),
+                    task_group=task_group,
+                )
+            )
+        await asyncio.Event().wait()

src/enapter/standalone/run.py

@@ -20,9 +20,8 @@ async def fixture_enapter_mqtt_client(mosquitto_container):
         port=int(ports[0]["HostPort"]),
     )
     async with asyncio.TaskGroup() as tg:
-        mqtt_client = enapter.mqtt.Client(tg, config)
-        yield mqtt_client
-        mqtt_client.cancel()
+        async with enapter.mqtt.Client(config, task_group=tg) as mqtt_client:
+            yield mqtt_client
 
 
 @pytest.fixture(scope="session", name="mosquitto_container")