"""Core topic/param wrapper over olo.core.v1.Core."""
from typing import Any, Callable, Dict, Iterator, List, Optional, Union, overload
import grpc
from olo_protos.core.v1 import core_pb2
from olo._convert import (
dict_to_struct,
param_value_to_python,
python_to_param_value,
seconds_to_duration,
struct_to_dict,
)
from olo.core._payloads import payload_from_proto, payload_to_proto
from olo.core.images import (
Image,
compressed_image_to_message,
compressed_message_to_frame,
image_message_to_frame,
image_to_message,
)
from olo.core.types import RobotNamespaces, ServiceInfo, TopicInfo
from olo.errors import wrap_rpc_error
_DEFAULT_WAIT_TIMEOUT_SEC = 5.0
_DEFAULT_RPC_TIMEOUT_SEC = 10.0
# Maps a ROS message type to a decoder turning the raw message dict into a typed object.
# get_latest() consults this by TopicInfo.msg_type so the caller's declared topic type drives
# the return shape (e.g. image topics yield an Image).
_DECODERS: Dict[str, Callable[[Dict[str, Any]], Any]] = {
"sensor_msgs/msg/Image": image_message_to_frame,
"sensor_msgs/msg/CompressedImage": compressed_message_to_frame,
}
# Mirror of _DECODERS for the publish path: maps a ROS message type to an encoder
# turning an Image into a ROS-shaped dict. publish() consults this by
# TopicInfo.msg_type so the declared topic type (not the Image) drives the output
# shape, keeping publish() symmetric with get_latest().
_ENCODERS: Dict[str, Callable[[Image], Dict[str, Any]]] = {
"sensor_msgs/msg/Image": image_to_message,
"sensor_msgs/msg/CompressedImage": compressed_image_to_message,
}
[docs]
class Subscription:
"""Context-managed sync iterator over a Core.Subscribe server stream."""
[docs]
def __init__(self, stream):
"""Wrap a Subscribe server stream."""
self._stream = stream
self._closed = False
def __iter__(self) -> Iterator[Dict[str, Any]]:
"""Iterate incoming messages."""
return self
def __next__(self) -> Dict[str, Any]:
"""Return the next subscribed message."""
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 payload_from_proto(response.message)
[docs]
def close(self) -> None:
"""Cancel the stream and mark it closed."""
if self._closed:
return
self._closed = True
cancel = getattr(self._stream, "cancel", None)
if callable(cancel):
cancel()
def __enter__(self) -> "Subscription":
"""Enter a context manager."""
return self
def __exit__(self, exc_type, exc, tb) -> None:
"""Close the subscription on context exit."""
self.close()
[docs]
class Publisher:
"""Lightweight handle for repeated Publish calls (stateless on the appliance)."""
[docs]
def __init__(self, core: "Core", topic: TopicInfo):
"""Capture the publish target."""
self._core = core
self.topic = topic
@overload
def publish(
self,
message: Dict[str, Any],
*,
timeout: Optional[float] = None,
) -> None: ...
@overload
def publish(
self,
message: Image,
*,
timeout: Optional[float] = None,
) -> None: ...
[docs]
def publish(
self,
message: Union[Dict[str, Any], Image],
*,
timeout: Optional[float] = None,
) -> None:
"""Publish a message on this topic."""
if isinstance(message, Image):
self._core.publish(self.topic, message, timeout=timeout)
else:
self._core.publish(self.topic, message, timeout=timeout)
[docs]
class CoreHandle:
"""Namespace-scoped handle for common core operations."""
[docs]
def __init__(self, core: "Core", namespace: str):
"""Bind to a core client and resolved robot namespace."""
self._core = core
self.namespace = namespace
[docs]
def topic(self, name: str, msg_type: str) -> TopicInfo:
"""Resolve a topic name against the robot namespace (ROS-style).
Leading ``/`` means absolute: the name is used as-is. Otherwise the name
is resolved relative to this handle's namespace.
"""
if name.startswith("/"):
return TopicInfo(name=name, msg_type=msg_type)
prefix = f"/{self.namespace}" if self.namespace else ""
return TopicInfo(name=f"{prefix}/{name}", msg_type=msg_type)
def _resolve(self, topic: TopicInfo) -> TopicInfo:
"""Resolve a :class:`TopicInfo` through this handle's namespace."""
return self.topic(topic.name, topic.msg_type)
[docs]
def node(self, name: str) -> str:
"""Build a namespace-prefixed node name."""
relative = name.lstrip("/")
if self.namespace:
return f"{self.namespace}/{relative}"
return relative
[docs]
def service(self, name: str) -> str:
"""Resolve a service name against the robot namespace (ROS-style)."""
if name.startswith("/"):
return name
prefix = f"/{self.namespace}" if self.namespace else ""
return f"{prefix}/{name.lstrip('/')}"
[docs]
def list_topics(self, *, timeout: float = _DEFAULT_RPC_TIMEOUT_SEC) -> List[TopicInfo]:
"""List topics visible through this handle using namespace-relative names."""
topics = self._core.list_topics(timeout=timeout)
if not self.namespace:
return [TopicInfo(name=topic.name.lstrip("/"), msg_type=topic.msg_type) for topic in topics]
prefix = f"/{self.namespace}/"
return [
TopicInfo(name=topic.name[len(prefix) :], msg_type=topic.msg_type)
for topic in topics
if topic.name.startswith(prefix)
]
[docs]
def get_latest(
self,
topic: TopicInfo,
*,
timeout: float = _DEFAULT_WAIT_TIMEOUT_SEC,
) -> Dict[str, Any]:
"""Wait for the next message on a topic resolved against this namespace."""
return self._core.get_latest(self._resolve(topic), timeout=timeout)
[docs]
def get_image(
self,
topic: TopicInfo,
*,
timeout: float = _DEFAULT_WAIT_TIMEOUT_SEC,
) -> Image:
"""Wait for the next image on a topic resolved against this namespace."""
return self._core.get_image(self._resolve(topic), timeout=timeout)
@overload
def publish(
self,
topic: TopicInfo,
message: Dict[str, Any],
*,
timeout: Optional[float] = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None: ...
@overload
def publish(
self,
topic: TopicInfo,
message: Image,
*,
timeout: Optional[float] = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None: ...
[docs]
def publish(
self,
topic: TopicInfo,
message: Union[Dict[str, Any], Image],
*,
timeout: Optional[float] = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None:
"""Publish a ROS message on a topic resolved against this namespace."""
resolved = self._resolve(topic)
if isinstance(message, Image):
self._core.publish(resolved, message, timeout=timeout)
else:
self._core.publish(resolved, message, timeout=timeout)
[docs]
def publisher(self, topic: TopicInfo) -> Publisher:
"""Create a publisher for a topic resolved against this namespace."""
return self._core.publisher(self._resolve(topic))
[docs]
def subscribe(
self,
topic: TopicInfo,
*,
timeout: Optional[float] = None,
) -> Subscription:
"""Subscribe to a topic resolved against this namespace."""
return self._core.subscribe(self._resolve(topic), timeout=timeout)
[docs]
def get_params(
self,
node: str,
names: List[str],
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> Dict[str, Any]:
"""Read ROS parameters from a namespace-relative node."""
return self._core.get_params(self.node(node), names, timeout=timeout)
[docs]
def set_params(
self,
node: str,
params: Dict[str, Any],
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None:
"""Write ROS parameters on a namespace-relative node."""
self._core.set_params(self.node(node), params, timeout=timeout)
[docs]
def list_services(
self,
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> List[ServiceInfo]:
"""List services visible through this handle using namespace-relative names."""
services = self._core.list_services(timeout=timeout)
if not self.namespace:
return [
ServiceInfo(name=service.name.lstrip("/"), service_type=service.service_type)
for service in services
]
prefix = f"/{self.namespace}/"
return [
ServiceInfo(
name=service.name[len(prefix) :],
service_type=service.service_type,
)
for service in services
if service.name.startswith(prefix)
]
[docs]
def call_service(
self,
service: ServiceInfo,
request: Optional[Dict[str, Any]] = None,
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> Dict[str, Any]:
"""Call a service resolved against this namespace."""
return self._core.call_service(
ServiceInfo(name=self.service(service.name), service_type=service.service_type),
request or {},
timeout=timeout,
)
[docs]
class Core:
[docs]
def __init__(self, session):
"""Bind to a channel session."""
self._session = session
def _call(self, stub_method, request, timeout: Optional[float]):
"""Invoke a stub method and map gRPC errors."""
try:
return stub_method(request, timeout=timeout)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
[docs]
def get_latest(
self,
topic: TopicInfo,
*,
timeout: float = _DEFAULT_WAIT_TIMEOUT_SEC,
) -> Dict[str, Any]:
"""
Wait for the next message on a topic and return it as a raw dict. It
resolves once a message is published (or the timeout elapses).
Use `getImage` to decode image topics into an `Image`.
"""
request = core_pb2.GetLatestRequest(
topic=topic.name,
message_type=topic.msg_type,
timeout=seconds_to_duration(timeout),
)
response = self._call(
self._session.core_stub.GetLatest,
request,
timeout=timeout,
)
return payload_from_proto(response.message)
[docs]
def get_image(
self,
topic: TopicInfo,
*,
timeout: float = _DEFAULT_WAIT_TIMEOUT_SEC,
) -> Image:
"""Wait for the next image on a topic and decode it into an :class:`Image`.
Handles both ``sensor_msgs/msg/Image`` and
``sensor_msgs/msg/CompressedImage`` transparently, driven by
``topic.msg_type``.
"""
decoder = _DECODERS.get(topic.msg_type)
if decoder is None:
raise ValueError(
f"get_image does not support topic type {topic.msg_type!r}; "
"expected sensor_msgs/msg/Image or sensor_msgs/msg/CompressedImage"
)
message = self.get_latest(topic, timeout=timeout)
return decoder(message)
@overload
def publish(
self,
topic: TopicInfo,
message: Dict[str, Any],
*,
timeout: Optional[float] = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None: ...
@overload
def publish(
self,
topic: TopicInfo,
message: Image,
*,
timeout: Optional[float] = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None: ...
[docs]
def publish(
self,
topic: TopicInfo,
message: Union[Dict[str, Any], Image],
*,
timeout: Optional[float] = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None:
"""Publish a ROS message on a topic.
``message`` is either a plain ROS message dict or an :class:`Image`.
When an :class:`Image` is given, it is encoded according to
``topic.msg_type`` (``sensor_msgs/msg/Image`` -> raw,
``sensor_msgs/msg/CompressedImage`` -> compressed), mirroring how
:meth:`get_latest` decodes those types into an :class:`Image`.
"""
if isinstance(message, Image):
encoder = _ENCODERS.get(topic.msg_type)
if encoder is None:
raise ValueError(
f"Cannot publish an Image to topic type {topic.msg_type!r}; "
"expected sensor_msgs/msg/Image or sensor_msgs/msg/CompressedImage"
)
message = encoder(message)
request = core_pb2.PublishRequest(
topic=topic.name,
message_type=topic.msg_type,
message=payload_to_proto(topic.msg_type, message),
)
self._call(self._session.core_stub.Publish, request, timeout=timeout)
[docs]
def publisher(self, topic: TopicInfo) -> Publisher:
"""Create a reusable publisher handle."""
return Publisher(self, topic)
[docs]
def list_topics(self, *, timeout: float = _DEFAULT_RPC_TIMEOUT_SEC) -> List[TopicInfo]:
"""List available ROS topics."""
response = self._call(
self._session.core_stub.ListTopics,
core_pb2.ListTopicsRequest(),
timeout=timeout,
)
return [TopicInfo(name=name, msg_type=msg_type) for name, msg_type in sorted(response.topics.items())]
[docs]
def list_robot_namespaces(self, *, timeout: float = _DEFAULT_RPC_TIMEOUT_SEC) -> RobotNamespaces:
"""List available robot namespaces."""
response = self._call(
self._session.core_stub.ListRobotNamespaces,
core_pb2.ListRobotNamespacesRequest(),
timeout=timeout,
)
return RobotNamespaces(has_global=response.has_global, namespaces=tuple(response.namespaces))
[docs]
def list_services(
self,
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> List[ServiceInfo]:
"""List available ROS services."""
response = self._call(
self._session.core_stub.ListServices,
core_pb2.ListServicesRequest(),
timeout=timeout,
)
return [
ServiceInfo(name=service.name, service_type=service.service_type)
for service in sorted(response.services, key=lambda item: item.name)
]
[docs]
def call_service(
self,
service: ServiceInfo,
request: Optional[Dict[str, Any]] = None,
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> Dict[str, Any]:
"""Call a ROS service and return the response as a plain dict."""
proto_request = core_pb2.CallServiceRequest(
service=service.name,
service_type=service.service_type,
request=dict_to_struct(request or {}),
timeout=seconds_to_duration(timeout),
)
response = self._call(
self._session.core_stub.CallService,
proto_request,
timeout=timeout,
)
return struct_to_dict(response.response)
[docs]
def handle(self, namespace: Optional[str] = None) -> CoreHandle:
"""Return a namespace-scoped core handle."""
return CoreHandle(self, self._session.resolve_namespace(namespace))
[docs]
def subscribe(
self,
topic: TopicInfo,
*,
timeout: Optional[float] = None,
) -> Subscription:
"""Subscribe to a topic stream."""
request = core_pb2.SubscribeRequest(topic=topic.name, message_type=topic.msg_type)
try:
stream = self._session.core_stub.Subscribe(request, timeout=timeout)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
return Subscription(stream)
[docs]
def get_params(
self,
node: str,
names: List[str],
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> Dict[str, Any]:
"""Read ROS parameters from a node.
Values come back as natives matching the ROS parameter type
(bool, int, float, str, bytes, or a homogeneous list). Unset
parameters are returned as None.
"""
request = core_pb2.GetParamsRequest(node=node, names=names)
response = self._call(self._session.core_stub.GetParams, request, timeout=timeout)
return {name: param_value_to_python(value) for name, value in response.params.items()}
[docs]
def set_params(
self,
node: str,
params: Dict[str, Any],
*,
timeout: float = _DEFAULT_RPC_TIMEOUT_SEC,
) -> None:
"""Write ROS parameters on a node.
The Python type of each value determines the ROS parameter type
(bool, int, float, str, bytes, or a homogeneous list thereof).
"""
proto_params = {name: python_to_param_value(value) for name, value in params.items()}
request = core_pb2.SetParamsRequest(node=node, params=proto_params)
self._call(self._session.core_stub.SetParams, request, timeout=timeout)
__all__ = [
"Core",
"CoreHandle",
"Image",
"Publisher",
"RobotNamespaces",
"ServiceInfo",
"Subscription",
"TopicInfo",
]