feat: add dynamic tests and docs

Author: gusye1234

drive_events/core.py

@@ -1,18 +1,21 @@
 import inspect
 import asyncio
 from typing import Callable, Optional, Union, Any, Tuple
-from .types import BaseEvent, EventFunction, EventGroup, EventInput
+from .types import (
+    BaseEvent,
+    EventFunction,
+    EventGroup,
+    EventInput,
+    _SpecialEventReturn,
+    ReturnBehavior,
+)
 from .broker import BaseBroker
 from .utils import (
     logger,
     string_to_md5_hash,
 )
 
 
-def goto_events(group_markers: list[BaseEvent], any_return: Any):
-    pass
-
-
 class EventEngineCls:
     def __init__(self, name="default", broker: Optional[BaseBroker] = None):
         self.name = name
@@ -23,6 +26,9 @@ def __init__(self, name="default", broker: Optional[BaseBroker] = None):
     def reset(self):
         self.__event_maps = {}
 
+    def get_event_from_id(self, event_id: str) -> Optional[BaseEvent]:
+        return self.__event_maps.get(event_id)
+
     def make_event(self, func: Union[EventFunction, BaseEvent]) -> BaseEvent:
         if isinstance(func, BaseEvent):
             self.__event_maps[func.id] = func
@@ -80,23 +86,38 @@ async def invoke_event(
         this_run_ctx = {}
         queue: list[Tuple[BaseEvent, EventInput]] = [(event, event_input)]
 
-        async def run_event(current_event, current_event_input):
+        async def run_event(current_event: BaseEvent, current_event_input: Any):
             result = await current_event.solo_run(current_event_input, global_ctx)
             this_run_ctx[current_event.id] = result
-            for cand_event in self.__event_maps.values():
-                cand_event_parents = cand_event.parent_groups
-                for group_hash, group in cand_event_parents.items():
-                    if current_event.id in group.events and all(
-                        [event_id in this_run_ctx for event_id in group.events]
-                    ):
-                        this_group_returns = {
-                            event_id: this_run_ctx[event_id]
-                            for event_id in group.events
-                        }
-                        build_input = EventInput(
-                            group_name=group.name, results=this_group_returns
+            if isinstance(result, _SpecialEventReturn):
+                if result.behavior == ReturnBehavior.GOTO:
+                    group_markers, any_return = result.returns
+                    for group_marker in group_markers:
+                        this_group_returns = {current_event.id: any_return}
+                        build_input_goto = EventInput(
+                            group_name="$goto",
+                            results=this_group_returns,
+                            behavior=ReturnBehavior.GOTO,
                         )
-                        queue.append((cand_event, build_input))
+                        queue.append((group_marker, build_input_goto))
+                elif result.behavior == ReturnBehavior.ABORT:
+                    return
+            else:
+                # dispath to events who listen
+                for cand_event in self.__event_maps.values():
+                    cand_event_parents = cand_event.parent_groups
+                    for group_hash, group in cand_event_parents.items():
+                        if current_event.id in group.events and all(
+                            [event_id in this_run_ctx for event_id in group.events]
+                        ):
+                            this_group_returns = {
+                                event_id: this_run_ctx[event_id]
+                                for event_id in group.events
+                            }
+                            build_input = EventInput(
+                                group_name=group.name, results=this_group_returns
+                            )
+                            queue.append((cand_event, build_input))
 
         while len(queue):
             this_batch_events = queue[:max_async_events] if max_async_events else queue

drive_events/dynamic.py

@@ -0,0 +1,16 @@
+from typing import Any
+from .types import (
+    BaseEvent,
+    _SpecialEventReturn,
+    ReturnBehavior,
+)
+
+
+def goto_events(group_markers: list[BaseEvent], any_return: Any) -> _SpecialEventReturn:
+    return _SpecialEventReturn(
+        behavior=ReturnBehavior.GOTO, returns=(group_markers, any_return)
+    )
+
+
+def abort_this():
+    return _SpecialEventReturn(behavior=ReturnBehavior.ABORT, returns=None)

drive_events/types.py

@@ -2,7 +2,7 @@
 from enum import Enum
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Callable, Any, Awaitable, Optional, TypeVar, Generic
+from typing import Any, Awaitable, Optional, Union, Callable
 
 from .utils import (
     string_to_md5_hash,
@@ -39,8 +39,22 @@ class EventInput(EventGroupInput):
     pass
 
 
+@dataclass
+class _SpecialEventReturn:
+    behavior: ReturnBehavior
+    returns: Any
+
+    def __post_init__(self):
+        if not isinstance(self.behavior, ReturnBehavior):
+            raise TypeError(
+                f"behavior must be a ReturnBehavior, not {type(self.behavior)}"
+            )
+
+
 # (group_event_results, global ctx set by user) -> result
-EventFunction = Callable[[Optional[EventInput], Optional[Any]], Awaitable[Any]]
+EventFunction = Callable[
+    [Optional[EventInput], Optional[Any]], Awaitable[Union[Any, _SpecialEventReturn]]
+]
 
 
 @dataclass

readme.md

@@ -52,7 +52,6 @@ from drive_events import EventInput, default_drive
 async def hello(event: EventInput, global_ctx):
     print("hello")
 
-
 @default_drive.listen_groups([hello])
 async def world(event: EventInput, global_ctx):
     print("world")
@@ -77,5 +76,129 @@ await default_drive.invoke_event(EVENT, EVENT_INPUT, GLOBAL_CTX)
 
 Check out [examples](./examples) for more user cases!
 
+### Multi-Recv
+
+`drive_events` allow an event to be triggered only when a group of events are produced:
+
+<details>
+<summary> code snippet</summary>
+
+```python
+@default_drive.make_event
+async def start(event: EventInput, global_ctx):
+    print("start")
+    
+@default_drive.listen_groups([start])
+async def hello(event: EventInput, global_ctx):
+    return 1
+
+
+@default_drive.listen_groups([start])
+async def world(event: EventInput, global_ctx):
+    return 2
+
+
+@default_drive.listen_groups([hello, world])
+async def adding(event: EventInput, global_ctx):
+    results = event.results
+    print("adding", hello, world)
+    return results[hello.id] + results[world.id]
+
+
+results = asyncio.run(default_drive.invoke_event(start))
+assert results[adding.id] == 3
+```
+</details>
+
+
+### Parallel
+
+`drive_events` is perfect for workflows that have many network IO that can be awaited in parallel. If two events are listened to the same group of events, then they will be triggered at the same time:
+
+<details>
+<summary> code snippet</summary>
+
+```python
+@default_drive.make_event
+async def start(event: EventInput, global_ctx):
+    print("start")
+
+@default_drive.listen_groups([start])
+async def hello(event: EventInput, global_ctx):
+    print(datetime.now(), "hello")
+    await asyncio.sleep(0.2)
+    print(datetime.now(), "hello done")
+
+
+@default_drive.listen_groups([start])
+async def world(event: EventInput, global_ctx):
+    print(datetime.now(), "world")
+    await asyncio.sleep(0.2)
+    print(datetime.now(), "world done")
+
+
+asyncio.run(default_drive.invoke_event(start))
+```
+
+</details>
+
+
+
+### Dynamic
+
+`drive_events` is dynamic. You can use `goto` and `abort` to change the workflow at runtime:
+
+<details>
+<summary> code snippet for abort</summary>
+
+```python
+from drive_events.dynamic import goto_events, abort_this
+
+@default_drive.make_event
+async def a(event: EventInput, global_ctx):
+    return abort_this()
+
+@default_drive.listen_groups([a])
+async def b(event: EventInput, global_ctx):
+    assert False, "should not be called"
+    
+asyncio.run(default_drive.invoke_event(a))
+```
+
+</details>
+
+<details>
+<summary> code snippet for goto</summary>
+
+```python
+call_a_count = 0
+
+@default_drive.make_event
+async def a(event: EventInput, global_ctx):
+    global call_a_count
+    if call_a_count == 0:
+        assert event is None
+    elif call_a_count == 1:
+        assert event.behavior == ReturnBehavior.GOTO
+        assert event.results == {b.id: 2}
+        return abort_this()
+    call_a_count += 1
+    return 1
+
+@default_drive.listen_groups([a])
+async def b(event: EventInput, global_ctx):
+    return goto_events([a], 2)
+
+@default_drive.listen_groups([b])
+async def c(event: EventInput, global_ctx):
+    assert False, "should not be called"
+    
+asyncio.run(default_drive.invoke_event(a))
+```
+
+</details>
+
+
+
 
 

tests/test_define.py

@@ -54,6 +54,15 @@ async def a(event: EventInput, global_ctx):
     assert isinstance(a, BaseEvent)
 
 
+@pytest.mark.asyncio
+async def test_correct_get_id():
+    @default_drive.make_event
+    async def a(event: EventInput, global_ctx):
+        return 1
+
+    assert default_drive.get_event_from_id(a.id) == a
+
+
 @pytest.mark.asyncio
 async def test_order():
     @default_drive.make_event

tests/test_dynamic_run.py

@@ -0,0 +1,59 @@
+import pytest
+from drive_events import default_drive, EventInput
+from drive_events.types import ReturnBehavior, _SpecialEventReturn
+from drive_events.dynamic import goto_events, abort_this
+
+
+class DeliberateExcepion(Exception):
+    pass
+
+
+def test_special_event_init():
+    with pytest.raises(TypeError):
+        _SpecialEventReturn("fool", 1)
+
+
+@pytest.mark.asyncio
+async def test_abort():
+    @default_drive.make_event
+    async def a(event: EventInput, global_ctx):
+        assert global_ctx == {"test_ctx": 1}
+        return abort_this()
+
+    @default_drive.listen_groups([a])
+    async def b(event: EventInput, global_ctx):
+        assert False, "should not be called"
+
+    result = await default_drive.invoke_event(a, None, {"test_ctx": 1})
+    print(result)
+
+
+@pytest.mark.asyncio
+async def test_goto():
+    call_a_count = 0
+
+    @default_drive.make_event
+    async def a(event: EventInput, global_ctx):
+        nonlocal call_a_count
+        if call_a_count == 0:
+            assert event is None
+        elif call_a_count == 1:
+            assert event.behavior == ReturnBehavior.GOTO
+            assert event.group_name == "$goto"
+            assert event.results == {b.id: 2}
+            return abort_this()
+        else:
+            raise ValueError("should not be called more than twice")
+        call_a_count += 1
+        return 1
+
+    @default_drive.listen_groups([a])
+    async def b(event: EventInput, global_ctx):
+        return goto_events([a], 2)
+
+    @default_drive.listen_groups([b])
+    async def c(event: EventInput, global_ctx):
+        assert False, "should not be called"
+
+    result = await default_drive.invoke_event(a, None, {"test_ctx": 1})
+    print(result)