39 lines
936 B
Python
39 lines
936 B
Python
from collections.abc import Iterator
|
|
from typing import Any, Callable, Generic, TypeVar, final
|
|
|
|
|
|
T = TypeVar("T")
|
|
U = TypeVar("U")
|
|
|
|
Op = Callable[["Stream[T]"], "Stream[U]"]
|
|
|
|
|
|
@final
|
|
class Stream(Generic[T]):
|
|
_extensions: dict[str, Callable[..., Any]] = {}
|
|
|
|
def __init__(self, it: Iterator[T]) -> None:
|
|
self._it = it
|
|
|
|
def __iter__(self) -> Iterator[T]:
|
|
return self._it
|
|
|
|
@classmethod
|
|
def extension(cls, name: str | None = None):
|
|
"""Register a function as Stream.<name>(...). First arg will be the stream."""
|
|
|
|
def deco(fn: Callable[..., Any]):
|
|
cls._extensions[name or fn.__name__] = fn
|
|
return fn
|
|
|
|
return deco
|
|
|
|
def __getattr__(self, attr: str):
|
|
fn = self._extensions.get(attr)
|
|
if fn is None:
|
|
raise AttributeError(attr)
|
|
|
|
def bound(*args, **kwargs):
|
|
return fn(self, *args, **kwargs)
|
|
|
|
return bound
|