:sparkles: feat(handler): support Interator while outcome Outputs

3 files changed, 19 insertions, 15 deletions

diff --git a/src/infini/core.py b/src/infini/core.py
index 4eaeb4f8..cbe90924 100644
--- a/src/infini/core.py
+++ b/src/infini/core.py
@@ -17,7 +17,7 @@ class Core:
             return
         for handled_stream in self.handle(pre_intercepted_stream):
             if handled_stream.is_empty():
-                return  # TODO 处理未找到 Handler
+                return
             yield self.intercept(self.generate(handled_stream))
 
     def pre_intercept(self, input: Input) -> Input | Output:
diff --git a/src/infini/handler.py b/src/infini/handler.py
index dd6c7fd7..5a4e1649 100644
--- a/src/infini/handler.py
+++ b/src/infini/handler.py
@@ -1,6 +1,6 @@
 from infini.input import Input
 from infini.output import Output
-from infini.typing import List, RouterType, Callable
+from infini.typing import List, Any, RouterType, Callable, Generator
 from infini.queue import EventQueue
 
 
@@ -8,21 +8,23 @@ class Handler:
     handlers: List[RouterType]
 
     def input(self, input: Input):
-        queue = self.match(input.get_plain_text())
-        if queue.is_empty():
+        if (queue := self.match(input.get_plain_text())).is_empty():
             yield Output.empty()
             return
         while not queue.is_empty():
-            output = queue.pop()(input)
-            # TODO 允许 Handler 函数为迭代器
-            # if isinstance(output, GeneratorType):
-            #     for o in output:
-            #         yield o
-            yield output
-            if output.block:
-                return
+            if isinstance(stream := queue.pop()(input), Generator):  # TODO 新增测试
+                for output in stream:
+                    yield output
+                    if output.block:
+                        return
+            else:
+                yield stream
+                if stream.block:
+                    return
 
-    def match(self, text: str) -> EventQueue[Callable[[Input], Output]]:
+    def match(
+        self, text: str
+    ) -> EventQueue[Callable[[Input], Output | Generator[Output, Any, None]]]:
         queue = EventQueue()
 
         for handler in self.handlers:
diff --git a/src/infini/typing.py b/src/infini/typing.py
index a1273f02..d9db1e61 100644
--- a/src/infini/typing.py
+++ b/src/infini/typing.py
@@ -6,7 +6,7 @@ from typing import (
     Callable as Callable,
     Literal as Literal,
     Sequence as Sequence,
-    overload as overload,
+    Generator as Generator,
     TypeVar,
     TypedDict,
     Union,
@@ -15,9 +15,11 @@ from types import ModuleType as ModuleType, GeneratorType as GeneratorType
 from . import router, input, output
 
 T = TypeVar("T")
+Stream = Union["input.Input", "output.Output"]
+OutputGenerator = Generator["output.Output", Any, None]
 
 
 class RouterType(TypedDict):
     priority: int
     router: router.Router
-    handler: Callable[["input.Input"], Union["input.Input", "output.Output"]]
+    handler: Callable[["input.Input"], Union[Stream, OutputGenerator]]