add test for middleware

README.md

@@ -45,6 +45,7 @@ pip install aiosaber
 
 ```python
 from aiosaber import *
+
 @task
 def add(self, num):
     for i in range(100000):
@@ -67,3 +68,47 @@ num_ch = Channel.values(*list(range(100)))
 f = my_flow(num_ch)
 asyncio.run(f.start())
 ```
+
+## Middleware example
+
+```python
+from aiosaber import *
+
+class NameBuilder(BaseBuilder):
+    def __call__(self, com, *args, **kwargs):
+        super().__call__(com, *args, **kwargs)
+        com.context['name'] = type(com).__name__ + str(id(com))
+
+class ClientProvider(BaseExecutor):
+    async def __call__(self, com, **kwargs):
+        if not context.context.get('client'):
+            context.context['client'] = 'client'
+        return await super().__call__(com, **kwargs)
+
+class Filter(BaseHandler):
+    async def __call__(self, com, get, put, **kwargs):
+        async def filter_put(data):
+            if data is END or data > 3:
+                await put(data)
+
+        return await super().__call__(com, get, filter_put, **kwargs)
+
+@task
+async def add(self, num):
+    print(self.context['name'])
+    print(context.context['client'])
+    return num + 1
+
+@flow
+def myflow(num_ch):
+    return num_ch | add | view
+
+context.context.update({
+    'builders': [NameBuilder],
+    'executors': [ClientProvider],
+    'handlers': [Filter]
+})
+f = myflow(Channel.values(1, 2, 3, 4, 5))
+context.context.clear()
+asyncio.run(f.start())
+```

aiosaber/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.0.1'
+__version__ = '0.0.1.1'
 
 import warnings
 
@@ -9,6 +9,8 @@
 # noinspection PyUnresolvedReferences
 from .flow import *
 # noinspection PyUnresolvedReferences
+from .middleware import BaseHandler, BaseExecutor, BaseBuilder
+# noinspection PyUnresolvedReferences
 from .operators import *
 # noinspection PyUnresolvedReferences
 from .task import *

aiosaber/component.py

@@ -6,7 +6,7 @@
 
 from aiosaber import context
 from aiosaber.channel import Consumer
-from aiosaber.plugins import ContextExecutor
+from aiosaber.middleware import ContextExecutor
 from aiosaber.utility.typings import ChannelOutput, Builder, Executor
 
 if TYPE_CHECKING:

aiosaber/flow.py

@@ -4,7 +4,7 @@
 from aiosaber import context
 from aiosaber.channel import Channel
 from aiosaber.component import Component
-from aiosaber.plugins import DaskExecutorProvider
+from aiosaber.middleware import DaskExecutorProvider
 from aiosaber.utility.typings import ChannelOutput
 from aiosaber.utility.utils import class_deco
 

aiosaber/utility/utils.py

@@ -44,9 +44,13 @@ def first_exception(fut: asyncio.Future):
     return_when = return_when or first_exception
 
     def _done_callback(fut: asyncio.Future):
-        q.task_done()
         sem.release()
         futures.remove(fut)
+        # TODO may over called
+        try:
+            q.task_done()
+        except Exception:
+            pass
         # set waiter done once
         if not waiter.done() and return_when(fut):
             # propagate all state to current coro

tests/test_middleware.py

@@ -0,0 +1,45 @@
+from aiosaber import *
+
+
+def test_middleware():
+    class NameBuilder(BaseBuilder):
+        def __call__(self, com, *args, **kwargs):
+            super().__call__(com, *args, **kwargs)
+            com.context['name'] = type(com).__name__ + str(id(com))
+
+    class ClientProvider(BaseExecutor):
+        async def __call__(self, com, **kwargs):
+            if not context.context.get('client'):
+                context.context['client'] = 'client'
+            return await super().__call__(com, **kwargs)
+
+    class Filter(BaseHandler):
+        async def __call__(self, com, get, put, **kwargs):
+            async def filter_put(data):
+                if data is END or data > 3:
+                    await put(data)
+
+            return await super().__call__(com, get, filter_put, **kwargs)
+
+    @task
+    async def add(self, num):
+        print(self.context['name'])
+        print(context.context['client'])
+        return num + 1
+
+    @flow
+    def myflow(num_ch):
+        return num_ch | add | view
+
+    context.context.update({
+        'builders': [NameBuilder],
+        'executors': [ClientProvider],
+        'handlers': [Filter]
+    })
+    f = myflow(Channel.values(1, 2, 3, 4, 5))
+    context.context.clear()
+    asyncio.run(f.start())
+
+
+if __name__ == "__main__":
+    test_middleware()