diff --git a/permchain/channels.py b/permchain/channels.py deleted file mode 100644 index 4812c1e90..000000000 --- a/permchain/channels.py +++ /dev/null @@ -1,426 +0,0 @@ -import json -from abc import ABC, abstractmethod -from contextlib import asynccontextmanager, contextmanager -from typing import ( - Any, - AsyncContextManager, - AsyncGenerator, - Callable, - FrozenSet, - Generator, - Generic, - Optional, - Sequence, - Type, - TypeVar, - Union, - cast, -) -from typing import ContextManager as ContextManagerType - -from typing_extensions import Self - -Value = TypeVar("Value") -Update = TypeVar("Update") - - -class EmptyChannelError(Exception): - pass - - -class InvalidUpdateError(Exception): - pass - - -class Channel(Generic[Value, Update], ABC): - @property - @abstractmethod - def ValueType(self) -> Any: - """The type of the value stored in the channel.""" - - @property - @abstractmethod - def UpdateType(self) -> Any: - """The type of the update received by the channel.""" - - @contextmanager - @abstractmethod - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - """Return a new identical channel, optionally initialized from a checkpoint.""" - - @asynccontextmanager - async def aempty( - self, checkpoint: Optional[str] = None - ) -> AsyncGenerator[Self, None]: - """Return a new identical channel, optionally initialized from a checkpoint.""" - with self.empty(checkpoint) as value: - yield value - - @abstractmethod - def update(self, values: Sequence[Update]) -> None: - ... - - @abstractmethod - def get(self) -> Value: - ... - - @abstractmethod - def checkpoint(self) -> str | None: - ... - - -class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]): - """Stores the result of applying a binary operator to the current value and each new value. - - ```python - import operator - - total = BinaryOperatorAggregate(int, operator.add) - ``` - """ - - def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]): - self.typ = typ - self.operator = operator - - @property - def ValueType(self) -> Type[Value]: - """The type of the value stored in the channel.""" - return self.typ - - @property - def UpdateType(self) -> Type[Value]: - """The type of the update received by the channel.""" - return self.typ - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ, self.operator) - if checkpoint is not None: - empty.value = json.loads(checkpoint) - try: - yield empty - finally: - try: - del empty.value - except AttributeError: - pass - - def update(self, values: Sequence[Value]) -> None: - if not hasattr(self, "value"): - self.value = values[0] - values = values[1:] - - for value in values: - self.value = self.operator(self.value, value) - - def get(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - return json.dumps(self.value) - - -class LastValue(Generic[Value], Channel[Value, Value]): - """Stores the last value received.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - - @property - def ValueType(self) -> Type[Value]: - """The type of the value stored in the channel.""" - return self.typ - - @property - def UpdateType(self) -> Type[Value]: - """The type of the update received by the channel.""" - return self.typ - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.value = json.loads(checkpoint) - try: - yield empty - finally: - try: - del empty.value - except AttributeError: - pass - - def update(self, values: Sequence[Value]) -> None: - if len(values) != 1: - raise InvalidUpdateError() - - self.value = values[-1] - - def get(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - return json.dumps(self.value) - - -class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): - """Stores all values received, resets in each step.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - - @property - def ValueType(self) -> Type[Sequence[Value]]: - """The type of the value stored in the channel.""" - return Sequence[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Any: - """The type of the update received by the channel.""" - return Union[self.typ, Sequence[self.typ]] # type: ignore[name-defined] - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.queue = tuple(json.loads(checkpoint)) - try: - yield empty - finally: - try: - del empty.queue - except AttributeError: - pass - - def update(self, values: Sequence[Value | Sequence[Value]]) -> None: - self.queue = tuple( - cast(Value, v) - for value in values - for v in ( - (value,) - if isinstance(value, self.typ) - else cast(Sequence[Value], value) - ) - ) - - def get(self) -> Sequence[Value]: - try: - return self.queue - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - return json.dumps(self.queue) - - -class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): - """Stores all unique values received, resets in each step.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - - @property - def ValueType(self) -> Type[Sequence[Value]]: - """The type of the value stored in the channel.""" - return Sequence[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Any: - """The type of the update received by the channel.""" - return Union[self.typ, Sequence[self.typ]] # type: ignore[name-defined] - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.queue = tuple(json.loads(checkpoint)) - try: - yield empty - finally: - try: - del empty.queue - except AttributeError: - pass - - def update(self, values: Sequence[Value | Sequence[Value]]) -> None: - self.queue = tuple( - set( - cast(Value, v) - for value in values - for v in ( - (value,) - if isinstance(value, self.typ) - else cast(Sequence[Value], value) - ) - ) - ) - - def get(self) -> Sequence[Value]: - try: - return self.queue - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - return json.dumps(self.queue) - - -class Set(Generic[Value], Channel[FrozenSet[Value], Value]): - """Stores all unique values received.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - self.set = set[Value]() - - @property - def ValueType(self) -> Type[FrozenSet[Value]]: - """The type of the value stored in the channel.""" - return FrozenSet[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Type[Value]: - """The type of the update received by the channel.""" - return self.typ - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.set = set(json.loads(checkpoint)) - try: - yield empty - finally: - pass - - def update(self, values: Sequence[Value]) -> None: - self.set.update(values) - - def get(self) -> FrozenSet[Value]: - try: - return frozenset(self.set) - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - return json.dumps(list(self.set)) - - -class Stream(Generic[Value], Channel[Sequence[Value], Value]): - """Stores all unique values received.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - self.set = list[Value]() - - @property - def ValueType(self) -> Any: - """The type of the value stored in the channel.""" - return Sequence[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Type[Value]: - """The type of the update received by the channel.""" - return self.typ - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.set = json.loads(checkpoint) - try: - yield empty - finally: - pass - - def update(self, values: Sequence[Value]) -> None: - self.set.extend(values) - - def get(self) -> Sequence[Value]: - try: - return tuple(self.set) - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - return json.dumps(self.set) - - -class ContextManager(Generic[Value], Channel[Value, None]): - value: Value - - def __init__( - self, - ctx: Optional[Callable[[], ContextManagerType[Value]]] = None, - actx: Optional[Callable[[], AsyncContextManager[Value]]] = None, - typ: Optional[Type[Value]] = None, - ) -> None: - if ctx is None and actx is None: - raise ValueError("Must provide either sync or async context manager.") - - self.typ = typ - self.ctx = ctx - self.actx = actx - - @property - def ValueType(self) -> Any: - """The type of the value stored in the channel.""" - return ( - self.typ - or (self.ctx if hasattr(self.ctx, "__enter__") else None) - or (self.actx if hasattr(self.actx, "__aenter__") else None) - or None - ) - - @property - def UpdateType(self) -> Type[None]: - """The type of the update received by the channel.""" - raise InvalidUpdateError() - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - if self.ctx is None: - raise ValueError("Cannot enter sync context manager.") - - empty = self.__class__(ctx=self.ctx, actx=self.actx, typ=self.typ) - # ContextManager doesn't have a checkpoint - ctx = self.ctx() - empty.value = ctx.__enter__() - try: - yield empty - finally: - ctx.__exit__(None, None, None) - - @asynccontextmanager - async def aempty( - self, checkpoint: Optional[str] = None - ) -> AsyncGenerator[Self, None]: - if self.actx is not None: - empty = self.__class__(ctx=self.ctx, actx=self.actx, typ=self.typ) - # ContextManager doesn't have a checkpoint - actx = self.actx() - empty.value = await actx.__aenter__() - try: - yield empty - finally: - await actx.__aexit__(None, None, None) - else: - with self.empty() as empty: - yield empty - - def update(self, values: Sequence[None]) -> None: - raise InvalidUpdateError() - - def get(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> None: - return None diff --git a/permchain/channels/__init__.py b/permchain/channels/__init__.py new file mode 100644 index 000000000..23bca3a44 --- /dev/null +++ b/permchain/channels/__init__.py @@ -0,0 +1,19 @@ +from permchain.channels.base import Channel, EmptyChannelError, InvalidUpdateError +from permchain.channels.inbox import Inbox, UniqueInbox +from permchain.channels.last_value import LastValue +from permchain.channels.binop import BinaryOperatorAggregate +from permchain.channels.stream import Set, Stream +from permchain.channels.context import ContextManager + +__all__ = [ + "Channel", + "EmptyChannelError", + "InvalidUpdateError", + "LastValue", + "Inbox", + "UniqueInbox", + "BinaryOperatorAggregate", + "Set", + "Stream", + "ContextManager", +] diff --git a/permchain/channels/base.py b/permchain/channels/base.py new file mode 100644 index 000000000..a7e62cd43 --- /dev/null +++ b/permchain/channels/base.py @@ -0,0 +1,61 @@ +from abc import ABC, abstractmethod +from contextlib import asynccontextmanager, contextmanager +from typing import ( + Any, + AsyncGenerator, + Generator, + Generic, + Optional, + Sequence, + TypeVar, +) + +from typing_extensions import Self + +Value = TypeVar("Value") +Update = TypeVar("Update") + + +class EmptyChannelError(Exception): + pass + + +class InvalidUpdateError(Exception): + pass + + +class Channel(Generic[Value, Update], ABC): + @property + @abstractmethod + def ValueType(self) -> Any: + """The type of the value stored in the channel.""" + + @property + @abstractmethod + def UpdateType(self) -> Any: + """The type of the update received by the channel.""" + + @contextmanager + @abstractmethod + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + """Return a new identical channel, optionally initialized from a checkpoint.""" + + @asynccontextmanager + async def aempty( + self, checkpoint: Optional[str] = None + ) -> AsyncGenerator[Self, None]: + """Return a new identical channel, optionally initialized from a checkpoint.""" + with self.empty(checkpoint) as value: + yield value + + @abstractmethod + def update(self, values: Sequence[Update]) -> None: + ... + + @abstractmethod + def get(self) -> Value: + ... + + @abstractmethod + def checkpoint(self) -> str | None: + ... diff --git a/permchain/channels/binop.py b/permchain/channels/binop.py new file mode 100644 index 000000000..2e703f048 --- /dev/null +++ b/permchain/channels/binop.py @@ -0,0 +1,62 @@ +import json +from contextlib import contextmanager +from typing import Callable, Generator, Generic, Optional, Sequence, Type + +from typing_extensions import Self + +from permchain.channels.base import Channel, EmptyChannelError, Value + + +class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]): + """Stores the result of applying a binary operator to the current value and each new value. + + ```python + import operator + + total = BinaryOperatorAggregate(int, operator.add) + ``` + """ + + def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]): + self.typ = typ + self.operator = operator + + @property + def ValueType(self) -> Type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ, self.operator) + if checkpoint is not None: + empty.value = json.loads(checkpoint) + try: + yield empty + finally: + try: + del empty.value + except AttributeError: + pass + + def update(self, values: Sequence[Value]) -> None: + if not hasattr(self, "value"): + self.value = values[0] + values = values[1:] + + for value in values: + self.value = self.operator(self.value, value) + + def get(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> str: + return json.dumps(self.value) diff --git a/permchain/channels/context.py b/permchain/channels/context.py new file mode 100644 index 000000000..1248dbc78 --- /dev/null +++ b/permchain/channels/context.py @@ -0,0 +1,97 @@ +from contextlib import asynccontextmanager, contextmanager +from typing import ( + Any, + AsyncContextManager, + AsyncGenerator, + Callable, + Generator, + Generic, + Optional, + Sequence, + Type, +) +from typing import ContextManager as ContextManagerType + +from typing_extensions import Self + +from permchain.channels.base import ( + Channel, + EmptyChannelError, + InvalidUpdateError, + Value, +) + + +class ContextManager(Generic[Value], Channel[Value, None]): + value: Value + + def __init__( + self, + ctx: Optional[Callable[[], ContextManagerType[Value]]] = None, + actx: Optional[Callable[[], AsyncContextManager[Value]]] = None, + typ: Optional[Type[Value]] = None, + ) -> None: + if ctx is None and actx is None: + raise ValueError("Must provide either sync or async context manager.") + + self.typ = typ + self.ctx = ctx + self.actx = actx + + @property + def ValueType(self) -> Any: + """The type of the value stored in the channel.""" + return ( + self.typ + or (self.ctx if hasattr(self.ctx, "__enter__") else None) + or (self.actx if hasattr(self.actx, "__aenter__") else None) + or None + ) + + @property + def UpdateType(self) -> Type[None]: + """The type of the update received by the channel.""" + raise InvalidUpdateError() + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + if self.ctx is None: + raise ValueError("Cannot enter sync context manager.") + + empty = self.__class__(ctx=self.ctx, actx=self.actx, typ=self.typ) + # ContextManager doesn't have a checkpoint + ctx = self.ctx() + empty.value = ctx.__enter__() + try: + yield empty + finally: + ctx.__exit__(None, None, None) + + @asynccontextmanager + async def aempty( + self, checkpoint: Optional[str] = None + ) -> AsyncGenerator[Self, None]: + if self.actx is not None: + empty = self.__class__(ctx=self.ctx, actx=self.actx, typ=self.typ) + # ContextManager doesn't have a checkpoint + actx = self.actx() + empty.value = await actx.__aenter__() + try: + yield empty + finally: + await actx.__aexit__(None, None, None) + else: + with self.empty() as empty: + yield empty + + def update(self, values: Sequence[None]) -> None: + raise InvalidUpdateError() + + def get(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> None: + return None diff --git a/permchain/channels/inbox.py b/permchain/channels/inbox.py new file mode 100644 index 000000000..afe9c90df --- /dev/null +++ b/permchain/channels/inbox.py @@ -0,0 +1,113 @@ +import json +from contextlib import contextmanager +from typing import Any, Generator, Generic, Optional, Sequence, Type, Union, cast + +from typing_extensions import Self + +from permchain.channels.base import ( + Channel, + EmptyChannelError, + Value, +) + + +class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): + """Stores all values received, resets in each step.""" + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + + @property + def ValueType(self) -> Type[Sequence[Value]]: + """The type of the value stored in the channel.""" + return Sequence[self.typ] # type: ignore[name-defined] + + @property + def UpdateType(self) -> Any: + """The type of the update received by the channel.""" + return Union[self.typ, Sequence[self.typ]] # type: ignore[name-defined] + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + empty.queue = tuple(json.loads(checkpoint)) + try: + yield empty + finally: + try: + del empty.queue + except AttributeError: + pass + + def update(self, values: Sequence[Value | Sequence[Value]]) -> None: + self.queue = tuple( + cast(Value, v) + for value in values + for v in ( + (value,) + if isinstance(value, self.typ) + else cast(Sequence[Value], value) + ) + ) + + def get(self) -> Sequence[Value]: + try: + return self.queue + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> str: + return json.dumps(self.queue) + + +class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): + """Stores all unique values received, resets in each step.""" + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + + @property + def ValueType(self) -> Type[Sequence[Value]]: + """The type of the value stored in the channel.""" + return Sequence[self.typ] # type: ignore[name-defined] + + @property + def UpdateType(self) -> Any: + """The type of the update received by the channel.""" + return Union[self.typ, Sequence[self.typ]] # type: ignore[name-defined] + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + empty.queue = tuple(json.loads(checkpoint)) + try: + yield empty + finally: + try: + del empty.queue + except AttributeError: + pass + + def update(self, values: Sequence[Value | Sequence[Value]]) -> None: + self.queue = tuple( + set( + cast(Value, v) + for value in values + for v in ( + (value,) + if isinstance(value, self.typ) + else cast(Sequence[Value], value) + ) + ) + ) + + def get(self) -> Sequence[Value]: + try: + return self.queue + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> str: + return json.dumps(self.queue) diff --git a/permchain/channels/last_value.py b/permchain/channels/last_value.py new file mode 100644 index 000000000..266e44ce0 --- /dev/null +++ b/permchain/channels/last_value.py @@ -0,0 +1,57 @@ +import json +from contextlib import contextmanager +from typing import Generator, Generic, Optional, Sequence, Type + +from typing_extensions import Self + +from permchain.channels.base import ( + Channel, + EmptyChannelError, + InvalidUpdateError, + Value, +) + + +class LastValue(Generic[Value], Channel[Value, Value]): + """Stores the last value received.""" + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + + @property + def ValueType(self) -> Type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + empty.value = json.loads(checkpoint) + try: + yield empty + finally: + try: + del empty.value + except AttributeError: + pass + + def update(self, values: Sequence[Value]) -> None: + if len(values) != 1: + raise InvalidUpdateError() + + self.value = values[-1] + + def get(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> str: + return json.dumps(self.value) diff --git a/permchain/channels/stream.py b/permchain/channels/stream.py new file mode 100644 index 000000000..d51b02b8d --- /dev/null +++ b/permchain/channels/stream.py @@ -0,0 +1,87 @@ +import json +from contextlib import contextmanager +from typing import Any, FrozenSet, Generator, Generic, Optional, Sequence, Type + +from typing_extensions import Self + +from permchain.channels.base import Channel, EmptyChannelError, Value + + +class Set(Generic[Value], Channel[FrozenSet[Value], Value]): + """Stores all unique values received.""" + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + self.set = set[Value]() + + @property + def ValueType(self) -> Type[FrozenSet[Value]]: + """The type of the value stored in the channel.""" + return FrozenSet[self.typ] # type: ignore[name-defined] + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + empty.set = set(json.loads(checkpoint)) + try: + yield empty + finally: + pass + + def update(self, values: Sequence[Value]) -> None: + self.set.update(values) + + def get(self) -> FrozenSet[Value]: + try: + return frozenset(self.set) + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> str: + return json.dumps(list(self.set)) + + +class Stream(Generic[Value], Channel[Sequence[Value], Value]): + """Stores all unique values received.""" + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + self.set = list[Value]() + + @property + def ValueType(self) -> Any: + """The type of the value stored in the channel.""" + return Sequence[self.typ] # type: ignore[name-defined] + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + empty.set = json.loads(checkpoint) + try: + yield empty + finally: + pass + + def update(self, values: Sequence[Value]) -> None: + self.set.extend(values) + + def get(self) -> Sequence[Value]: + try: + return tuple(self.set) + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> str: + return json.dumps(self.set) diff --git a/permchain/pregel.py b/permchain/pregel.py index d15f7ef9b..4934ebf2b 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -44,7 +44,7 @@ from langchain.schema.runnable.config import ( ) from langchain.schema.runnable.utils import ConfigurableFieldSpec -from permchain.channels import Channel, EmptyChannelError +from permchain.channels.base import Channel, EmptyChannelError logger = logging.getLogger(__name__)