Source code for olo.spatial.tree_subscription

"""Client wrapper for Spatial.SubscribeTree server streams."""

from collections.abc import Iterator
from dataclasses import dataclass

import grpc

from olo.errors import wrap_rpc_error
from olo.spatial._convert import transform_from_proto
from olo.spatial._tree import build_frame_tree
from olo.spatial.types import FrameTree, Transform




@dataclass(frozen=True, slots=True)
class TreeDelta:
    """Raw tree change batch from a SubscribeTree stream."""

    added: tuple[Transform, ...]
    updated: tuple[Transform, ...]
    root_frames: tuple[str, ...]
    base_frame: str
    snapshot: bool





class TreeSubscription:
    """Context-managed sync iterator over a Spatial.SubscribeTree stream.

    Default iteration yields a materialized :class:`FrameTree` after each
    server message. Use :meth:`changes` for the raw delta view. Only one
    consumer should iterate a subscription at a time.
    """



    def __init__(self, stream, *, namespace: str = "") -> None:
        """Wrap a SubscribeTree server stream."""
        self._stream = stream
        self._namespace = namespace
        self._closed = False
        self._edges: dict[str, Transform] = {}
        self._root_frames: set[str] = set()
        self._base_frame = ""


    def _apply(self, response) -> TreeDelta:
        pre_update_children = set(self._edges)
        added: list[Transform] = []
        updated: list[Transform] = []

        if response.is_snapshot:
            self._edges.clear()
            self._root_frames = set(response.root_frames)
            self._base_frame = response.base_frame
            for message in response.upserts:
                transform = transform_from_proto(message)
                self._edges[transform.child_frame_id] = transform
                added.append(transform)
            return TreeDelta(
                added=tuple(added),
                updated=(),
                root_frames=tuple(response.root_frames),
                base_frame=response.base_frame,
                snapshot=True,
            )

        if response.root_frames:
            self._root_frames = set(response.root_frames)
        if response.base_frame:
            self._base_frame = response.base_frame

        for message in response.upserts:
            transform = transform_from_proto(message)
            if transform.child_frame_id in pre_update_children:
                updated.append(transform)
            else:
                added.append(transform)
            self._edges[transform.child_frame_id] = transform

        return TreeDelta(
            added=tuple(added),
            updated=tuple(updated),
            root_frames=tuple(response.root_frames),
            base_frame=response.base_frame,
            snapshot=False,
        )

    def _materialized_tree(self) -> FrameTree:
        return build_frame_tree(self._edges.values(), self._root_frames, self._namespace)

    def _read_delta(self) -> TreeDelta:
        if self._closed:
            raise StopIteration
        try:
            response = next(self._stream)
        except StopIteration:
            self._closed = True
            raise
        except grpc.RpcError as exc:
            self._closed = True
            raise wrap_rpc_error(exc) from exc
        return self._apply(response)

    def __iter__(self) -> Iterator[FrameTree]:
        return self

    def __next__(self) -> FrameTree:
        self._read_delta()
        return self._materialized_tree()



    def changes(self) -> Iterator[TreeDelta]:
        """Iterate raw tree deltas instead of materialized trees."""
        while True:
            yield self._read_delta()




    def close(self) -> None:
        """Cancel the stream and mark it closed."""
        if self._closed:
            return
        self._closed = True
        self._stream.cancel()


    def __enter__(self) -> "TreeSubscription":
        """Enter a context manager."""
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        """Close the subscription on context exit."""
        self.close()