pymake

commit 793fe8c5fccee20a5da4ccea27e8423fe66f453f
parent e79043a791765498026ab48c3f7f7de80e1f5c6a
Author: gracefu <81774659+gracefuu@users.noreply.github.com>
Date:   Wed, 16 Apr 2025 22:51:11 +0800

Automatically wrap non-coroutines in coroutines when detaching

Diffstat:
M
examples/examples.py
 | 
4
++--
M
make/__init__.py
 | 
35
++++++++++++++++++++++-------------

2 files changed, 24 insertions(+), 15 deletions(-)
diff --git a/examples/examples.py b/examples/examples.py
@@ -86,10 +86,10 @@ async def main():
 
         # Note that `build(...)` will wait for all detached jobs to complete before returning.
         # You may instead use `build.build(...)`, which does not wait for detached jobs.
-        # You should ensure `build.wait()` is called eventually so detached jobs can complete.
+        # You should ensure `detach.wait()` is called eventually so detached jobs can complete.
         detach(build.build(_eg_rec(2345)))
         await build.build(_eg_rec(3456))
-        await build.wait()
+        await detach.wait()
 
 
 if __name__ == "__main__":
diff --git a/make/__init__.py b/make/__init__.py
@@ -341,10 +341,26 @@ class Detach:
     def __init__(self):
         self._background_tasks = set()
 
+    @staticmethod
+    async def _coro(fut):
+        return await fut
+
     def __call__(self, *args, **kwargs):
-        task = asyncio.create_task(*args, **kwargs)
+        awaitable, *args = args
+        if not hasattr(awaitable, "send"):
+            # the awaitable is NOT a coroutine, wrap it into one
+            awaitable = Detach._coro(awaitable)
+        task = asyncio.create_task(awaitable, *args, **kwargs)
         self._background_tasks.add(task)
         task.add_done_callback(self._background_tasks.discard)
+        return task
+
+    async def wait(self):
+        while self._background_tasks:
+            await asyncio.gather(*self._background_tasks)
+            self._background_tasks = set(
+                (t for t in self._background_tasks if not t.done())
+            )
 
 
 detach = Detach()
@@ -361,12 +377,8 @@ class SuspendingScheduler:
         self.done = set()
         self.waits = dict()
 
-    async def wait(self):
-        while detach._background_tasks:
-            await asyncio.gather(*detach._background_tasks)
-
-    async def build(self, *tasks: Task):
-        return await asyncio.gather(*(self.fetch_once(task) for task in tasks))
+    def build(self, *tasks: Task):
+        return asyncio.gather(*(self.fetch_once(task) for task in tasks))
 
     async def fetch_once(self, task: Task):
         task_key = task.task_key
@@ -407,14 +419,11 @@ class Build:
 
     async def __call__(self, *tasks: Task):
         result = await self.build(*tasks)
-        await self.wait()
+        await detach.wait()
         return result
 
-    async def wait(self):
-        return await self._scheduler.wait()
-
-    async def build(self, *tasks: Task):
-        return await self._scheduler.build(*tasks)
+    def build(self, *tasks: Task):
+        return self._scheduler.build(*tasks)
 
     def __enter__(self):
         self._store.__enter__()