streams

qas_flow/funcs.py

@@ -0,0 +1,76 @@
+from typing import Callable
+from .stream import Stream, T, U
+
+
+@Stream.extension()
+def map(stream: Stream[T], fn: Callable[[T], U]) -> Stream[U]:
+    def gen():
+        for x in stream:
+            yield fn(x)
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def filter(stream: Stream[T], pred: Callable[[T], bool]) -> Stream[T]:
+    def gen():
+        for x in stream:
+            if pred(x):
+                yield x
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def take(stream: Stream[T], n: int) -> Stream[T]:
+    def gen():
+        c = 0
+        for x in stream:
+            if c < n:
+                c += 1
+                yield x
+            else:
+                return
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def skip(stream: Stream[T], n: int) -> Stream[T]:
+    def gen():
+        c = 0
+        for x in stream:
+            c += 1
+            if c > n:
+                yield x
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def batch(stream: Stream[T], n: int) -> Stream[list[T]]:
+    def gen():
+        ls: list[T] = []
+        for x in stream:
+            ls.append(x)
+            if len(ls) == n:
+                yield ls
+                ls = []
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def enumerate(stream: Stream[T]) -> Stream[tuple[int, T]]:
+    def gen():
+        idx = 0
+        for x in stream:
+            yield (idx, x)
+            idx += 1
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def collect(stream: Stream[T]) -> list[T]:
+    return [v for v in stream]