From 4cf9d2306903a43e0b325d1188d8f14a6f2067dd Mon Sep 17 00:00:00 2001
From: QuantumGhost <obelisk.reg+git@gmail.com>
Date: Mon, 26 May 2025 14:22:57 +0800
Subject: [PATCH] feat(api): Add some utility functions for working with
 generators

---
 api/libs/gen_utils.py | 53 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)
 create mode 100644 api/libs/gen_utils.py

diff --git a/api/libs/gen_utils.py b/api/libs/gen_utils.py
new file mode 100644
index 0000000000..e457176187
--- /dev/null
+++ b/api/libs/gen_utils.py
@@ -0,0 +1,53 @@
+"""Utility functions for working with generators."""
+
+import logging
+from collections.abc import Callable, Generator
+from inspect import isgenerator
+from typing import TypeVar
+
+_YieldT = TypeVar("_YieldT")
+_YieldR = TypeVar("_YieldR")
+
+_T = TypeVar("_T")
+
+
+def inspect(gen_or_normal: _T, logger: logging.Logger) -> _T:
+    if not isgenerator(gen_or_normal):
+        return gen_or_normal
+
+    def wrapper():
+        for item in gen_or_normal:
+            logger.info(
+                "received generator item, type=%s, value=%s",
+                type(item),
+                item,
+            )
+            yield item
+
+    return wrapper()
+
+
+def map_(
+    gen: Generator[_YieldT, None, None],
+    mapper: Callable[[_YieldT], _YieldR],
+) -> Generator[_YieldR, None, None]:
+    for item in gen:
+        yield mapper(item)
+
+
+def filter_(
+    gen: Generator[_YieldT, None, None],
+    mapper: Callable[[_YieldT], bool],
+) -> Generator[_YieldT, None, None]:
+    for item in gen:
+        if mapper(item):
+            yield item
+
+
+def wrap(
+    gen: Generator[_YieldT, None, None],
+    funcs: list[Callable[[Generator[_YieldT, None, None]], Generator[_YieldT, None, None]]],
+) -> Generator[_YieldT, None, None]:
+    for f in funcs:
+        gen = f(gen)
+    return gen