Source code for olo.spatial

"""Spatial maths and TF wrapper over olo.spatial.v1.Spatial."""

from collections.abc import Iterator
from datetime import datetime

import grpc
from olo_protos.spatial.v1 import spatial_pb2

from olo._convert import datetime_to_timestamp, seconds_to_duration
from olo.errors import wrap_rpc_error
from olo.spatial._convert import transform_from_proto, transform_stamped_to_proto
from olo.spatial._timeouts import _DEFAULT_TIMEOUT_SEC, _LOOKUP_RPC_OVERHEAD_SEC
from olo.spatial._tree import build_frame_tree
from olo.spatial.frames import RobotFrames
from olo.spatial.tree_subscription import TreeDelta, TreeSubscription
from olo.spatial.types import (
    FrameInfo,
    FrameTree,
    FrameTreeNode,
    Point,
    Pose,
    Quaternion,
    Transform,
    Twist,
    Vector,
    wrap_error,
)
from olo.utils.common import normalize_robot_namespace




class TransformSubscription:
    """Context-managed sync iterator over a Spatial.SubscribeTransform stream."""



    def __init__(self, stream):
        """Wrap a SubscribeTransform server stream."""
        self._stream = stream
        self._closed = False


    def __iter__(self) -> Iterator[Transform]:
        """Iterate incoming transforms."""
        return self

    def __next__(self) -> Transform:
        """Return the next transform update."""
        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 transform_from_proto(response.transform)



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


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

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





class Spatial:
    """TF tree access via the Spatial gRPC service."""



    def __init__(self, session) -> None:
        """Bind to a channel session."""
        self._session = session




    def lookup(
        self,
        target_frame: str,
        source_frame: str,
        *,
        at: datetime | None = None,
        timeout: float = _DEFAULT_TIMEOUT_SEC,
        robot_namespace: str | None = None,
    ) -> Transform:
        """Look up a transform between two TF frames.

        ``target_frame`` is the frame poses are expressed in; ``source_frame``
        is the frame being looked up. The returned :class:`Transform` has
        ``frame_id=target_frame`` and ``child_frame_id=source_frame``.

        When ``robot_namespace`` is set, bare frame names are resolved
        server-side for that robot (see :class:`RobotFrames`).

        ``at`` selects a historical transform; when omitted the latest value is
        returned. ``timeout`` bounds how long to wait for the transform to
        become available.
        """
        request = spatial_pb2.LookupTransformRequest(
            target_frame=target_frame,
            source_frame=source_frame,
            timeout=seconds_to_duration(timeout),
            robot_namespace=normalize_robot_namespace(robot_namespace),
        )
        if at is not None:
            request.at.CopyFrom(datetime_to_timestamp(at))
        try:
            response = self._session.spatial_stub.LookupTransform(
                request,
                timeout=timeout + _LOOKUP_RPC_OVERHEAD_SEC,
            )
        except grpc.RpcError as exc:
            raise wrap_rpc_error(exc) from exc
        return transform_from_proto(response.transform)




    def tree(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> FrameTree:
        """Return the full TF tree across all robots and namespaces."""
        response = self._get_tree(None, timeout=timeout)
        return build_frame_tree(
            (transform_from_proto(edge) for edge in response.edges),
            response.root_frames,
        )




    def handle(self, namespace: str | None = None) -> RobotFrames:
        """Return a namespace-scoped TF frame handle for a robot.

        ``None`` auto-resolves the unambiguous appliance default. ``""`` or
        ``"global"`` select the global namespace without discovery; bare frame
        names are not prefixed server-side. For the full TF tree across all
        robots, use :meth:`tree` instead.
        """
        return RobotFrames(self, self._session.resolve_namespace(namespace))


    def _get_tree(
        self,
        namespace: str | None,
        *,
        include_ancestors: bool = False,
        timeout: float = _DEFAULT_TIMEOUT_SEC,
    ) -> spatial_pb2.GetTreeResponse:
        """Fetch a tree snapshot response."""
        request = spatial_pb2.GetTreeRequest(include_ancestors=include_ancestors)
        if namespace is not None:
            request.robot_namespace = normalize_robot_namespace(namespace)
        try:
            return self._session.spatial_stub.GetTree(request, timeout=timeout)
        except grpc.RpcError as exc:
            raise wrap_rpc_error(exc) from exc

    def _subscribe_tree(
        self,
        namespace: str | None,
        *,
        include_ancestors: bool = False,
        timeout: float | None = None,
    ) -> TreeSubscription:
        """Subscribe to TF tree updates for a full or scoped tree."""
        request = spatial_pb2.SubscribeTreeRequest(include_ancestors=include_ancestors)
        if namespace is not None:
            request.robot_namespace = normalize_robot_namespace(namespace)
        try:
            stream = self._session.spatial_stub.SubscribeTree(
                request,
                timeout=timeout,
            )
        except grpc.RpcError as exc:
            raise wrap_rpc_error(exc) from exc
        return TreeSubscription(stream, namespace=normalize_robot_namespace(namespace))



    def subscribe(
        self,
        target_frame: str,
        source_frame: str,
        *,
        timeout: float | None = None,
        robot_namespace: str | None = None,
    ) -> TransformSubscription:
        """Subscribe to transform updates between two frames.

        ``target_frame`` and ``source_frame`` follow the same convention as
        :meth:`lookup`.

        When ``robot_namespace`` is set, bare frame names are resolved
        server-side for that robot (see :class:`RobotFrames`).

        ``timeout`` is a deadline for the whole stream, not the wait for the
        next update: the subscription is terminated with a timeout error once
        it elapses. Leave it ``None`` (the default) for a long-lived
        subscription and close it via the context manager or ``close()``.
        """
        request = spatial_pb2.SubscribeTransformRequest(
            target_frame=target_frame,
            source_frame=source_frame,
            robot_namespace=normalize_robot_namespace(robot_namespace),
        )
        try:
            stream = self._session.spatial_stub.SubscribeTransform(
                request,
                timeout=timeout,
            )
        except grpc.RpcError as exc:
            raise wrap_rpc_error(exc) from exc
        return TransformSubscription(stream)




    def subscribe_tree(
        self,
        *,
        timeout: float | None = None,
    ) -> TreeSubscription:
        """Subscribe to TF tree updates across the full tree.

        Each yield is a materialized :class:`FrameTree` with current edge poses.
        Use :meth:`TreeSubscription.changes` on the returned subscription for
        raw deltas.

        ``timeout`` is a deadline for the whole stream, not the wait for the
        next update. Leave it ``None`` (the default) for a long-lived
        subscription and close it via the context manager or ``close()``.
        """
        return self._subscribe_tree(None, timeout=timeout)




    def publish_static(
        self,
        transform: Transform,
        *,
        timeout: float = _DEFAULT_TIMEOUT_SEC,
        robot_namespace: str | None = None,
    ) -> None:
        """Publish a static transform to the TF tree.

        The transform is broadcast on ``/tf_static``. ``frame_id`` and
        ``child_frame_id`` must be set. When ``stamp`` is omitted the server
        stamps the transform with its current time.

        When ``robot_namespace`` is set, bare frame ids are resolved
        server-side for that robot (see :class:`RobotFrames`).
        """
        request = spatial_pb2.PublishStaticTransformRequest(
            transform=transform_stamped_to_proto(transform),
            robot_namespace=normalize_robot_namespace(robot_namespace),
        )
        try:
            self._session.spatial_stub.PublishStaticTransform(request, timeout=timeout)
        except grpc.RpcError as exc:
            raise wrap_rpc_error(exc) from exc




__all__ = [
    "FrameTree",
    "FrameInfo",
    "FrameTreeNode",
    "Point",
    "Pose",
    "Quaternion",
    "RobotFrames",
    "Spatial",
    "Transform",
    "Twist",
    "Vector",
    "TransformSubscription",
    "TreeDelta",
    "TreeSubscription",
    "wrap_error",
]