From 184a065b76b711ffcbf240a2c0d41f97fd08f58a Mon Sep 17 00:00:00 2001 From: LJ Date: Wed, 2 Apr 2025 19:23:39 -0700 Subject: [PATCH] Make `FlowUpdater` provide async context manager. --- python/cocoindex/flow.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/cocoindex/flow.py b/python/cocoindex/flow.py index 9e330408..031f4187 100644 --- a/python/cocoindex/flow.py +++ b/python/cocoindex/flow.py @@ -390,6 +390,13 @@ def __exit__(self, exc_type, exc_value, traceback): self.abort() asyncio.run(self.wait()) + async def __aenter__(self) -> FlowLiveUpdater: + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + self.abort() + await self.wait() + async def wait(self) -> None: """ Wait for the live updater to finish. @@ -448,7 +455,7 @@ def name(self) -> str: """ return self._lazy_engine_flow().name() - async def update(self): + async def update(self) -> _engine.IndexUpdateInfo: """ Update the index defined by the flow. Once the function returns, the indice is fresh up to the moment when the function is called.