General support for generators

executor/engine/job/base.py

@@ -10,7 +10,7 @@
 
 from .utils import (
     JobStatusAttr, InvalidStateError, JobStatusType,
-    ExecutorError, valid_job_statuses
+    ExecutorError, valid_job_statuses, GeneratorWrapper
 )
 from .condition import Condition, AfterOthers, AllSatisfied
 from ..middle.capture import CaptureOut
@@ -275,7 +275,10 @@ async def wait_and_run(self):
                 self.status = "running"
                 try:
                     res = await self.run()
-                    await self.on_done(res)
+                    if not isinstance(res, GeneratorWrapper):
+                        await self.on_done(res)
+                    else:
+                        self.future.set_result(res)
                     return res
                 except Exception as e:
                     await self.on_failed(e)
@@ -285,7 +288,22 @@ async def wait_and_run(self):
 
     async def run(self):
         """Run the job."""
-        pass
+        if inspect.isgeneratorfunction(self.func) or inspect.isasyncgenfunction(self.func):  # noqa: E501
+            return await self.run_generator()
+        else:
+            return await self.run_function()
+
+    async def run_function(self):  # pragma: no cover
+        """Run the job as a function."""
+        msg = f"{type(self).__name__} does not implement " \
+              "run_function method."
+        raise NotImplementedError(msg)
+
+    async def run_generator(self):  # pragma: no cover
+        """Run the job as a generator."""
+        msg = f"{type(self).__name__} does not implement " \
+              "run_generator method."
+        raise NotImplementedError(msg)
 
     async def rerun(self, check_status: bool = True):
         """Rerun the job."""
@@ -351,9 +369,11 @@ def clear_context(self):
 
     def result(self) -> T.Any:
         """Get the result of the job."""
-        if self.status != "done":
-            raise InvalidStateError(self, ['done'])
-        return self.future.result()
+        res = self.future.result()
+        if not isinstance(res, GeneratorWrapper):
+            if self.status != "done":
+                raise InvalidStateError(self, ['done'])
+        return res
 
     def exception(self):
         """Get the exception of the job."""

executor/engine/job/dask.py

@@ -53,7 +53,7 @@ def release_resource(self) -> bool:
                 True
             )
 
-    async def run(self):
+    async def run_function(self):
         """Run job with Dask."""
         client = self.engine.dask_client
         func = functools.partial(self.func, **self.kwargs)

executor/engine/job/local.py

@@ -2,7 +2,7 @@
 
 
 class LocalJob(Job):
-    async def run(self):
+    async def run_function(self):
         """Run job in local thread."""
         res = self.func(*self.args, **self.kwargs)
         return res

executor/engine/job/process.py

@@ -1,6 +1,5 @@
 import asyncio
 import functools
-import inspect
 
 from loky.process_executor import ProcessPoolExecutor
 
@@ -43,18 +42,21 @@ def release_resource(self) -> bool:
                 True
             )
 
-    async def run(self):
+    async def run_function(self):
         """Run job in process pool."""
         func = functools.partial(self.func, *self.args, **self.kwargs)
-        if (inspect.isgeneratorfunction(self.func) or inspect.isasyncgenfunction(self.func)):  # noqa: E501
-            self._executor = ProcessPoolExecutor(
-                1, initializer=_gen_initializer, initargs=(func,))
-            result = GeneratorWrapper(self)
-        else:
-            self._executor = ProcessPoolExecutor(1)
-            loop = asyncio.get_running_loop()
-            fut = loop.run_in_executor(self._executor, func)
-            result = await fut
+        self._executor = ProcessPoolExecutor(1)
+        loop = asyncio.get_running_loop()
+        fut = loop.run_in_executor(self._executor, func)
+        result = await fut
+        return result
+
+    async def run_generator(self):
+        """Run job as a generator."""
+        func = functools.partial(self.func, *self.args, **self.kwargs)
+        self._executor = ProcessPoolExecutor(
+            1, initializer=_gen_initializer, initargs=(func,))
+        result = GeneratorWrapper(self)
         return result
 
     async def cancel(self):

executor/engine/job/thread.py

@@ -40,7 +40,7 @@ def release_resource(self) -> bool:
                 True
             )
 
-    async def run(self):
+    async def run_function(self):
         """Run job in thread pool."""
         self._executor = ThreadPoolExecutor(1)
         loop = asyncio.get_running_loop()

executor/engine/job/utils.py

@@ -66,11 +66,33 @@ def __iter__(self):
         return self
 
     def __next__(self) -> _T:
-        return self._job._executor.submit(_gen_next).result()
+        try:
+            return self._job._executor.submit(_gen_next).result()
+        except Exception as e:
+            engine = self._job.engine
+            if engine is None:
+                loop = asyncio.get_event_loop()
+            else:
+                loop = engine.loop
+            if isinstance(e, StopIteration):
+                cor = self._job.on_done(self)
+            else:
+                cor = self._job.on_failed(e)
+            fut = asyncio.run_coroutine_threadsafe(cor, loop)
+            fut.result()
+            raise e
 
     def __aiter__(self):
         return self
 
     async def __anext__(self) -> _T:
-        fut = self._job._executor.submit(_gen_anext)
-        return (await asyncio.wrap_future(fut))
+        try:
+            fut = self._job._executor.submit(_gen_anext)
+            res = await asyncio.wrap_future(fut)
+            return res
+        except Exception as e:
+            if isinstance(e, StopAsyncIteration):
+                await self._job.on_done(self)
+            else:
+                await self._job.on_failed(e)
+            raise e

executor/engine/launcher/core.py

@@ -127,6 +127,10 @@ def _fetch_result(job: 'Job') -> T.Any:
         else:
             return job.result()
 
+    def __call__(
+            self, *args: T.Any, **kwargs: T.Any) -> T.Any:  # pragma: no cover
+        raise NotImplementedError("Subclasses must implement __call__")
+
 
 class SyncLauncher(LauncherBase):
 

tests/test_job.py

@@ -188,24 +188,57 @@ def add(a, b):
 
 @pytest.mark.asyncio
 async def test_generator():
-    def gen():
-        for i in range(10):
-            yield i
-
-    async def gen_async(n):
-        for i in range(n):
-            yield i
-
     with Engine() as engine:
+        def gen():
+            for i in range(10):
+                yield i
+
         job = ProcessJob(gen)
         await engine.submit_async(job)
         await job.join()
-        assert list(job.result()) == list(range(10))
+        assert job.status == "running"
+        g = job.result()
+        assert list(g) == list(range(10))
+        assert job.status == "done"
+
+        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()
         res = []
         async for i in job.result():
+            assert job.status == "running"
             res.append(i)
+        assert job.status == "done"
         assert res == list(range(10))
+
+        def gen_error():
+            for i in range(2):
+                print(i)
+                yield i
+            raise ValueError("error")
+
+        job = ProcessJob(gen_error)
+        await engine.submit_async(job)
+        await job.join()
+        with pytest.raises(ValueError):
+            for i in job.result():
+                assert job.status == "running"
+        assert job.status == "failed"
+
+        async def gen_error():
+            for i in range(2):
+                print(i)
+                yield i
+            raise ValueError("error")
+
+        job = ProcessJob(gen_error)
+        await engine.submit_async(job)
+        await job.join()
+        with pytest.raises(ValueError):
+            async for i in job.result():
+                assert job.status == "running"
+        assert job.status == "failed"