feat(api): Add some utility functions for working with generators

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