loopix-sdk 2.30.0__py3-none-any.whl

loopix/sandbox_async/filesystem/watch_handle.py

@@ -0,0 +1,97 @@
+import asyncio
+import inspect
+
+from typing import Any, AsyncGenerator, Awaitable, Callable, Optional
+
+from loopix.envd.rpc import ahandle_rpc_exception_with_health
+from loopix.envd.filesystem.filesystem_pb2 import WatchDirResponse
+from loopix.sandbox.filesystem.filesystem import map_entry_info
+from loopix.sandbox.filesystem.watch_handle import FilesystemEvent, map_event_type
+from loopix.sandbox_async.utils import OutputHandler
+
+
+class AsyncWatchHandle:
+    """
+    Handle for watching a directory in the sandbox filesystem.
+
+    Use `.stop()` to stop watching the directory.
+    """
+
+    def __init__(
+        self,
+        events: AsyncGenerator[WatchDirResponse, Any],
+        on_event: OutputHandler[FilesystemEvent],
+        on_exit: Optional[OutputHandler[Optional[Exception]]] = None,
+        check_health: Optional[Callable[[], Awaitable[Optional[bool]]]] = None,
+    ):
+        self._events = events
+        self._on_event = on_event
+        self._on_exit = on_exit
+        self._check_health = check_health
+
+        self._wait = asyncio.create_task(self._handle_events())
+
+    async def stop(self):
+        """
+        Stop watching the directory.
+        """
+        self._wait.cancel()
+        await asyncio.wait([self._wait])
+        try:
+            await self._events.aclose()
+        except Exception:
+            pass
+
+    async def _iterate_events(self):
+        try:
+            async for event in self._events:
+                if event.HasField("filesystem"):
+                    event_type = map_event_type(event.filesystem.type)
+                    if event_type:
+                        yield FilesystemEvent(
+                            name=event.filesystem.name,
+                            type=event_type,
+                            entry=(
+                                map_entry_info(event.filesystem.entry)
+                                if event.filesystem.HasField("entry")
+                                else None
+                            ),
+                        )
+        except Exception as e:
+            raise await ahandle_rpc_exception_with_health(e, self._check_health)
+
+    async def _call_on_exit(self, error: Optional[Exception]):
+        if self._on_exit is None:
+            return
+        try:
+            cb = self._on_exit(error)
+            if inspect.isawaitable(cb):
+                await cb
+        except Exception:
+            # `on_exit` is the terminal callback; an error it raises has nowhere
+            # to propagate in this background task, so it's swallowed to avoid an
+            # "Task exception was never retrieved" warning. A `CancelledError`
+            # (a `BaseException`) is intentionally not caught here.
+            pass
+
+    async def _handle_events(self):
+        error: Optional[Exception] = None
+        try:
+            async for event in self._iterate_events():
+                cb = self._on_event(event)
+                if inspect.isawaitable(cb):
+                    await cb
+        except asyncio.CancelledError:
+            # `stop()` cancels this task to end the watch. Treat it as a clean,
+            # user-initiated end: fire `on_exit` (with no error), then propagate
+            # the cancellation so the task still finishes as cancelled.
+            await self._call_on_exit(None)
+            raise
+        except Exception as e:
+            error = e
+
+        # `on_exit` fires exactly once when the watch ends — with the error when
+        # the stream failed, or with `None` on a clean end. This matches the JS
+        # SDK, which calls `onExit()` after the loop completes and `onExit(err)`
+        # on error.
+        await self._call_on_exit(error)