coco-framework 0.1.0__py3-none-any.whl

coco_graph/channels/last_value.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Generic
+
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.errors import (
+    EmptyChannelError,
+    ErrorCode,
+    InvalidUpdateError,
+    create_error_message,
+)
+
+__all__ = ("LastValue", "LastValueAfterFinish")
+
+
+class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
+    """Stores the last value received, can receive at most one value per step."""
+
+    __slots__ = ("value",)
+
+    value: Value | Any
+
+    def __init__(self, typ: Any, key: str = "") -> None:
+        super().__init__(typ, key)
+        self.value = MISSING
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, LastValue)
+
+    @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
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.key)
+        empty.value = self.value
+        return empty
+
+    def from_checkpoint(self, checkpoint: Value) -> Self:
+        empty = self.__class__(self.typ, self.key)
+        if checkpoint is not MISSING:
+            empty.value = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        if len(values) == 0:
+            return False
+        if len(values) != 1:
+            msg = create_error_message(
+                message=f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values.",
+                error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
+            )
+            raise InvalidUpdateError(msg)
+
+        self.value = values[-1]
+        return True
+
+    def get(self) -> Value:
+        if self.value is MISSING:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING
+
+    def checkpoint(self) -> Value:
+        return self.value
+
+
+class LastValueAfterFinish(
+    Generic[Value], BaseChannel[Value, Value, tuple[Value, bool]]
+):
+    """Stores the last value received, but only made available after finish().
+    Once made available, clears the value."""
+
+    __slots__ = ("value", "finished")
+
+    value: Value | Any
+    finished: bool
+
+    def __init__(self, typ: Any, key: str = "") -> None:
+        super().__init__(typ, key)
+        self.value = MISSING
+        self.finished = False
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, LastValueAfterFinish)
+
+    @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
+
+    def checkpoint(self) -> tuple[Value | Any, bool] | Any:
+        if self.value is MISSING:
+            return MISSING
+        return (self.value, self.finished)
+
+    def from_checkpoint(self, checkpoint: tuple[Value | Any, bool] | Any) -> Self:
+        empty = self.__class__(self.typ)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            empty.value, empty.finished = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value | Any]) -> bool:
+        if len(values) == 0:
+            return False
+
+        self.finished = False
+        self.value = values[-1]
+        return True
+
+    def consume(self) -> bool:
+        if self.finished:
+            self.finished = False
+            self.value = MISSING
+            return True
+
+        return False
+
+    def finish(self) -> bool:
+        if not self.finished and self.value is not MISSING:
+            self.finished = True
+            return True
+        else:
+            return False
+
+    def get(self) -> Value:
+        if self.value is MISSING or not self.finished:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING and self.finished

coco_graph/channels/named_barrier_value.py

@@ -0,0 +1,167 @@
+from collections.abc import Sequence
+from typing import Generic
+
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.errors import EmptyChannelError, InvalidUpdateError
+
+__all__ = ("NamedBarrierValue", "NamedBarrierValueAfterFinish")
+
+
+class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
+    """A channel that waits until all named values are received before making the value available."""
+
+    __slots__ = ("names", "seen")
+
+    names: set[Value]
+    seen: set[Value]
+
+    def __init__(self, typ: type[Value], names: set[Value]) -> None:
+        super().__init__(typ)
+        self.names = names
+        self.seen: set[str] = set()
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, NamedBarrierValue) and value.names == self.names
+
+    @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
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.names)
+        empty.key = self.key
+        empty.seen = self.seen.copy()
+        return empty
+
+    def checkpoint(self) -> set[Value]:
+        return self.seen
+
+    def from_checkpoint(self, checkpoint: set[Value]) -> Self:
+        empty = self.__class__(self.typ, self.names)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            empty.seen = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        updated = False
+        for value in values:
+            if value in self.names:
+                if value not in self.seen:
+                    self.seen.add(value)
+                    updated = True
+            else:
+                raise InvalidUpdateError(
+                    f"At key '{self.key}': Value {value} not in {self.names}"
+                )
+        return updated
+
+    def get(self) -> Value:
+        if self.seen != self.names:
+            raise EmptyChannelError()
+        return None
+
+    def is_available(self) -> bool:
+        return self.seen == self.names
+
+    def consume(self) -> bool:
+        if self.seen == self.names:
+            self.seen = set()
+            return True
+        return False
+
+
+class NamedBarrierValueAfterFinish(
+    Generic[Value], BaseChannel[Value, Value, set[Value]]
+):
+    """A channel that waits until all named values are received before making the value ready to be made available. It is only made available after finish() is called."""
+
+    __slots__ = ("names", "seen", "finished")
+
+    names: set[Value]
+    seen: set[Value]
+
+    def __init__(self, typ: type[Value], names: set[Value]) -> None:
+        super().__init__(typ)
+        self.names = names
+        self.seen: set[str] = set()
+        self.finished = False
+
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, NamedBarrierValueAfterFinish)
+            and value.names == self.names
+        )
+
+    @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
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.names)
+        empty.key = self.key
+        empty.seen = self.seen.copy()
+        empty.finished = self.finished
+        return empty
+
+    def checkpoint(self) -> tuple[set[Value], bool]:
+        return (self.seen, self.finished)
+
+    def from_checkpoint(self, checkpoint: tuple[set[Value], bool]) -> Self:
+        empty = self.__class__(self.typ, self.names)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            empty.seen, empty.finished = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        updated = False
+        for value in values:
+            if value in self.names:
+                if value not in self.seen:
+                    self.seen.add(value)
+                    updated = True
+            else:
+                raise InvalidUpdateError(
+                    f"At key '{self.key}': Value {value} not in {self.names}"
+                )
+        return updated
+
+    def get(self) -> Value:
+        if not self.finished or self.seen != self.names:
+            raise EmptyChannelError()
+        return None
+
+    def is_available(self) -> bool:
+        return self.finished and self.seen == self.names
+
+    def consume(self) -> bool:
+        if self.finished and self.seen == self.names:
+            self.finished = False
+            self.seen = set()
+            return True
+        return False
+
+    def finish(self) -> bool:
+        if not self.finished and self.seen == self.names:
+            self.finished = True
+            return True
+        else:
+            return False

coco_graph/channels/topic.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from typing import Any, Generic
+
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.errors import EmptyChannelError
+
+__all__ = ("Topic",)
+
+
+def _flatten(values: Sequence[Value | list[Value]]) -> Iterator[Value]:
+    for value in values:
+        if isinstance(value, list):
+            yield from value
+        else:
+            yield value
+
+
+class Topic(
+    Generic[Value],
+    BaseChannel[Sequence[Value], Value | list[Value], list[Value]],
+):
+    """A configurable PubSub Topic.
+
+    Args:
+        typ: The type of the value stored in the channel.
+        accumulate: Whether to accumulate values across steps. If `False`, the channel will be emptied after each step.
+    """
+
+    __slots__ = ("values", "accumulate")
+
+    def __init__(self, typ: type[Value], accumulate: bool = False) -> None:
+        super().__init__(typ)
+        # attrs
+        self.accumulate = accumulate
+        # state
+        self.values = list[Value]()
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, Topic) and value.accumulate == self.accumulate
+
+    @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) -> Any:
+        """The type of the update received by the channel."""
+        return self.typ | list[self.typ]  # type: ignore[name-defined]
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.accumulate)
+        empty.key = self.key
+        empty.values = self.values.copy()
+        return empty
+
+    def checkpoint(self) -> list[Value]:
+        return self.values
+
+    def from_checkpoint(self, checkpoint: list[Value]) -> Self:
+        empty = self.__class__(self.typ, self.accumulate)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            if isinstance(checkpoint, tuple):
+                # backwards compatibility
+                empty.values = checkpoint[1]
+            else:
+                empty.values = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value | list[Value]]) -> bool:
+        updated = False
+        if not self.accumulate:
+            updated = bool(self.values)
+            self.values = list[Value]()
+        if flat_values := tuple(_flatten(values)):
+            updated = True
+            self.values.extend(flat_values)
+        return updated
+
+    def get(self) -> Sequence[Value]:
+        if self.values:
+            return list(self.values)
+        else:
+            raise EmptyChannelError
+
+    def is_available(self) -> bool:
+        return bool(self.values)

coco_graph/channels/untracked_value.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Generic
+
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.errors import EmptyChannelError, InvalidUpdateError
+
+__all__ = ("UntrackedValue",)
+
+
+class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
+    """Stores the last value received, never checkpointed."""
+
+    __slots__ = ("value", "guard")
+
+    guard: bool
+    value: Value | Any
+
+    def __init__(self, typ: type[Value], guard: bool = True) -> None:
+        super().__init__(typ)
+        self.guard = guard
+        self.value = MISSING
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, UntrackedValue) and value.guard == self.guard
+
+    @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
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.guard)
+        empty.key = self.key
+        empty.value = self.value
+        return empty
+
+    def checkpoint(self) -> Value | Any:
+        return MISSING
+
+    def from_checkpoint(self, checkpoint: Value) -> Self:
+        empty = self.__class__(self.typ, self.guard)
+        empty.key = self.key
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        if len(values) == 0:
+            return False
+        if len(values) != 1 and self.guard:
+            raise InvalidUpdateError(
+                f"At key '{self.key}': UntrackedValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
+            )
+
+        self.value = values[-1]
+        return True
+
+    def get(self) -> Value:
+        if self.value is MISSING:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING