diff options
| -rw-r--r-- | src/infini/core.py | 2 | ||||
| -rw-r--r-- | src/infini/handler.py | 26 | ||||
| -rw-r--r-- | src/infini/typing.py | 6 |
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]] |