fix for support generator

Author: Nanguage

README.md

@@ -77,11 +77,6 @@ from executor.engine import Engine, ProcessJob
 def add(a, b):
     return a + b
 
-async def stream():
-    for i in range(5):
-      await asyncio.sleep(0.5)
-      yield i
-
 with Engine() as engine:
     # job1 and job2 will be executed in parallel
     job1 = ProcessJob(add, args=(1, 2))
@@ -92,12 +87,6 @@ with Engine() as engine:
     engine.wait_job(job3)  # wait for job3 done
     print(job3.result())  # 10
 
-    # generator
-    job4 = ProcessJob(stream)
-    # do not do engine.wait because the generator job's future is done only when StopIteration
-    await engine.submit_async(job4)
-    async for x in job3.result():
-      print(x)
 ```
 
 Async mode example:
@@ -111,6 +100,11 @@ engine = Engine()
 def add(a, b):
     return a + b
 
+async def stream():
+    for i in range(5):
+      await asyncio.sleep(0.5)
+      yield i
+
 async def main():
     job1 = ProcessJob(add, args=(1, 2))
     job2 = ProcessJob(add, args=(job1.future, 4))
@@ -119,6 +113,13 @@ async def main():
     print(job1.result())  # 3
     print(job2.result())  # 7
 
+    # generator
+    job3 = ProcessJob(stream)
+    # do not do engine.wait because the generator job's future is done only when StopIteration
+    await engine.submit_async(job3)
+    async for x in job3.result():
+        print(x)
+
 asyncio.run(main())
 # or just `await main()` in jupyter environment
 ```

docs/getting-started.md

@@ -319,6 +319,48 @@ with Engine() as engine:
     engine.wait()
 ```
 
+### Generator support
+
+`executor.engine` supports generator job, which is a special job that returns a generator.
+
+```python
+import asyncio
+from executor.engine import Engine, ProcessJob
+
+engine = Engine()
+
+def gen():
+    for i in range(10):
+        yield i
+
+async def async_gen():
+    for i in range(10):
+        await asyncio.sleep(0.5)
+        yield i
+
+async def main():
+    job = ProcessJob(gen)
+    await engine.submit_async(job)
+    await job.wait_until_status("running")
+    for i in job.result():
+        print(i)
+
+    job = ProcessJob(async_gen)
+    await engine.submit_async(job)
+    await job.wait_until_status("running")
+    async for i in job.result():
+        print(i)
+
+asyncio.run(main())
+```
+
+!!! info
+    `LocalJob`, `ThreadJob`, `ProcessJob`, `DaskJob` support generator job.
+
+!!! warning
+    Do not use `engine.wait()` to wait the generator job done,
+    because the generator job's future is done only when the generator is exhausted.
+
 ## Engine
 
 `executor.engine` provides a `Engine` class for managing jobs.

executor/engine/__init__.py

@@ -1,7 +1,7 @@
 from .core import Engine, EngineSetting
 from .job import LocalJob, ThreadJob, ProcessJob
 
-__version__ = '0.2.5'
+__version__ = '0.2.6'
 
 __all__ = [
     'Engine', 'EngineSetting',

executor/engine/job/local.py

@@ -1,4 +1,5 @@
 from .base import Job
+from .utils import GeneratorWrapper
 
 
 class LocalJob(Job):
@@ -9,5 +10,4 @@ async def run_function(self):
 
     async def run_generator(self):
         """Run job as a generator."""
-        res = self.func(*self.args, **self.kwargs)
-        return res
+        return GeneratorWrapper(self)

executor/engine/job/utils.py

@@ -74,17 +74,25 @@ class GeneratorWrapper(T.Generic[_T]):
     def __init__(self, job: "Job", fut: T.Optional[Future] = None):
         self._job = job
         self._fut = fut
+        self._local_res = None
 
     def __iter__(self):
         return self
 
     def __next__(self) -> _T:
         try:
-            return self._job._executor.submit(_gen_next, self._fut).result()
+            if self._job._executor is not None:
+                return self._job._executor.submit(
+                    _gen_next, self._fut).result()
+            else:
+                if self._local_res is None:
+                    self._local_res = self._job.func(
+                        *self._job.args, **self._job.kwargs)
+                return next(self._local_res)
         except Exception as e:
             engine = self._job.engine
             if engine is None:
-                loop = asyncio.get_event_loop()
+                loop = asyncio.get_event_loop()  # pragma: no cover
             else:
                 loop = engine.loop
             if isinstance(e, StopIteration):
@@ -100,9 +108,15 @@ def __aiter__(self):
 
     async def __anext__(self) -> _T:
         try:
-            fut = self._job._executor.submit(_gen_anext, self._fut)
-            res = await asyncio.wrap_future(fut)
-            return res
+            if self._job._executor is not None:
+                fut = self._job._executor.submit(_gen_anext, self._fut)
+                res = await asyncio.wrap_future(fut)
+                return res
+            else:
+                if self._local_res is None:
+                    self._local_res = self._job.func(
+                        *self._job.args, **self._job.kwargs)
+                return await self._local_res.__anext__()
         except Exception as e:
             if isinstance(e, StopAsyncIteration):
                 await self._job.on_done(self)

tests/test_dask_job.py

@@ -76,3 +76,22 @@ async def main():
         client.close()
 
     asyncio.run(main())
+
+
+@pytest.mark.asyncio
+async def test_dask_generator():
+    with Engine() as engine:
+        async def gen():
+            for i in range(10):
+                yield i
+
+        job = DaskJob(gen)
+        await engine.submit_async(job)
+        await job.wait_until_status("running")
+        assert job.status == "running"
+        g = job.result()
+        i = 0
+        async for x in g:
+            assert x == i
+            i += 1
+        assert job.status == "done"

tests/test_job.py

@@ -195,26 +195,70 @@ def gen():
 
         job = ProcessJob(gen)
         await engine.submit_async(job)
-        await job.join()
+        await job.wait_until_status("running")
         assert job.status == "running"
         g = job.result()
         assert list(g) == list(range(10))
         assert job.status == "done"
 
+        job = ThreadJob(gen)
+        await engine.submit_async(job)
+        await job.wait_until_status("running")
+        assert job.status == "running"
+        g = job.result()
+        assert list(g) == list(range(10))
+        assert job.status == "done"
+
+        job = LocalJob(gen)
+        await engine.submit_async(job)
+        await job.wait_until_status("running")
+        assert job.status == "running"
+        g = job.result()
+        assert list(g) == list(range(10))
+        assert job.status == "done"
+
+
+@pytest.mark.asyncio
+async def test_generator_async():
+    with Engine() as engine:
         async def gen_async(n):
             for i in range(n):
                 yield i
 
         job = ProcessJob(gen_async, (10,))
         await engine.submit_async(job)
-        await job.join()
+        await job.wait_until_status("running")
+        res = []
+        async for i in job.result():
+            assert job.status == "running"
+            res.append(i)
+        assert job.status == "done"
+        assert res == list(range(10))
+
+        job = ThreadJob(gen_async, (10,))
+        await engine.submit_async(job)
+        await job.wait_until_status("running")
         res = []
         async for i in job.result():
             assert job.status == "running"
             res.append(i)
         assert job.status == "done"
         assert res == list(range(10))
 
+        job = LocalJob(gen_async, (10,))
+        await engine.submit_async(job)
+        await job.wait_until_status("running")
+        res = []
+        async for i in job.result():
+            assert job.status == "running"
+            res.append(i)
+        assert job.status == "done"
+        assert res == list(range(10))
+
+
+@pytest.mark.asyncio
+async def test_generator_error():
+    with Engine() as engine:
         def gen_error():
             for i in range(2):
                 print(i)
@@ -223,7 +267,7 @@ def gen_error():
 
         job = ProcessJob(gen_error)
         await engine.submit_async(job)
-        await job.join()
+        await job.wait_until_status("running")
         with pytest.raises(ValueError):
             for i in job.result():
                 assert job.status == "running"
@@ -237,7 +281,7 @@ async def gen_error():
 
         job = ProcessJob(gen_error)
         await engine.submit_async(job)
-        await job.join()
+        await job.wait_until_status("running")
         with pytest.raises(ValueError):
             async for i in job.result():
                 assert job.status == "running"