"""Appliance-local archive (saved images, videos, and rosbags)."""
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Optional
import grpc
from google.protobuf import timestamp_pb2
from olo_protos.archive.v1 import archive_pb2
from olo.archive.types import (
BagInfo,
BagSessionTerminalEvent,
BagSummary,
BagTopicInfo,
CaptureImageResult,
ImageSummary,
StopBagPlaybackResult,
StopBagRecordingResult,
StopVideoCaptureResult,
VideoSummary,
)
from olo.errors import wrap_rpc_error
from olo.utils.common import normalize_robot_namespace
if TYPE_CHECKING:
from olo._channel import ChannelSession
_DEFAULT_TIMEOUT_SEC = 30.0
_DEFAULT_CAPTURE_TIMEOUT_SEC = 20.0
_DEFAULT_WATCH_TIMEOUT_SEC = 300.0
def _timestamp_to_datetime(value: timestamp_pb2.Timestamp | None) -> datetime | None:
if value is None or (value.seconds == 0 and value.nanos == 0):
return None
return value.ToDatetime().replace(tzinfo=timezone.utc)
def _image_summary_from_proto(summary: archive_pb2.ImageSummary) -> ImageSummary:
return ImageSummary(
filename=summary.filename,
file_path=summary.file_path,
file_size=int(summary.file_size),
mime_type=summary.mime_type,
created_at=_timestamp_to_datetime(summary.created_at),
modified_at=_timestamp_to_datetime(summary.modified_at),
)
def _video_summary_from_proto(summary: archive_pb2.VideoSummary) -> VideoSummary:
return VideoSummary(
filename=summary.filename,
file_path=summary.file_path,
file_size=int(summary.file_size),
created_at=_timestamp_to_datetime(summary.created_at),
modified_at=_timestamp_to_datetime(summary.modified_at),
recording_id=summary.recording_id,
topic=summary.topic,
status=summary.status,
)
def _topic_from_proto(topic: archive_pb2.BagTopicInfo) -> BagTopicInfo:
return BagTopicInfo(
name=topic.name,
type=topic.type,
message_count=int(topic.message_count),
)
def _bag_summary_from_proto(summary: archive_pb2.BagSummary) -> BagSummary:
return BagSummary(
name=summary.name,
total_bytes=int(summary.total_bytes),
file_count=int(summary.file_count),
created_at=_timestamp_to_datetime(summary.created_at),
updated_at=_timestamp_to_datetime(summary.updated_at),
duration_seconds=float(summary.duration_seconds),
message_count=int(summary.message_count),
storage_format=summary.storage_format,
)
def _bag_from_proto(bag: archive_pb2.BagInfo) -> BagInfo:
if not bag.HasField("summary"):
raise ValueError("Malformed bag response")
return BagInfo(
summary=_bag_summary_from_proto(bag.summary),
topics=tuple(_topic_from_proto(topic) for topic in bag.topics),
compression=bag.compression,
)
def _terminal_from_event(event: archive_pb2.BagSessionEvent) -> BagSessionTerminalEvent:
if event.HasField("completed"):
return BagSessionTerminalEvent(
type="completed",
exit_code=int(event.completed.exit_code),
total_bytes=int(event.completed.total_bytes),
)
if event.HasField("failed"):
return BagSessionTerminalEvent(
type="failed",
exit_code=int(event.failed.exit_code),
error_message=event.failed.error_message,
)
raise ValueError("Session watch ended without terminal event")
[docs]
class VideoSession:
"""Handle for an active video capture session."""
[docs]
def __init__(self, client: "Archive", recording_id: str, filename: str, file_path: str) -> None:
self._client = client
self.id = recording_id
self.filename = filename
self.file_path = file_path
self._stopped = False
def stop(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> StopVideoCaptureResult:
result = self._client.stop_video_capture(self.id, timeout=timeout)
self._stopped = True
return result
def __enter__(self) -> "VideoSession":
return self
def __exit__(self, exc_type, exc, tb) -> None:
if not self._stopped:
self.stop()
[docs]
class BagSession:
"""Handle for an active bag recording session."""
[docs]
def __init__(self, client: "Archive", recording_id: str, filename: str, output_path: str) -> None:
self._client = client
self.id = recording_id
self.filename = filename
self.output_path = output_path
self._stopped = False
def stop(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> StopBagRecordingResult:
result = self._client.stop_bag_recording(self.id, timeout=timeout)
self._stopped = True
return result
def watch(self, *, timeout: float = _DEFAULT_WATCH_TIMEOUT_SEC) -> BagSessionTerminalEvent:
return self._client.watch_bag_session(self.id, timeout=timeout)
def __enter__(self) -> "BagSession":
return self
def __exit__(self, exc_type, exc, tb) -> None:
if not self._stopped:
self.stop()
[docs]
class PlaybackSession:
"""Handle for an active bag playback session."""
[docs]
def __init__(self, client: "Archive", playback_id: str, name: str) -> None:
self._client = client
self.id = playback_id
self.name = name
self._stopped = False
def stop(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> StopBagPlaybackResult:
result = self._client.stop_bag_playback(self.id, timeout=timeout)
self._stopped = True
return result
def pause(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> None:
self._client.pause_bag_playback(self.id, timeout=timeout)
def resume(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> None:
self._client.resume_bag_playback(self.id, timeout=timeout)
def seek(self, offset_seconds: float, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> None:
self._client.seek_bag_playback(self.id, offset_seconds, timeout=timeout)
def watch(self, *, timeout: float = _DEFAULT_WATCH_TIMEOUT_SEC) -> BagSessionTerminalEvent:
return self._client.watch_bag_session(self.id, timeout=timeout)
def __enter__(self) -> "PlaybackSession":
return self
def __exit__(self, exc_type, exc, tb) -> None:
if not self._stopped:
self.stop()
[docs]
class Archive:
"""Sync wrapper for appliance-local archive control."""
[docs]
def __init__(self, session: "ChannelSession") -> None:
self._session = session
def _call(self, stub_method, request, timeout: float):
try:
return stub_method(request, timeout=timeout)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
def capture_image(
self,
*,
topic: str,
message_type: str,
filename: str = "",
timeout_ms: int = 0,
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> CaptureImageResult:
response = self._call(
self._session.archive_stub.CaptureImage,
archive_pb2.CaptureImageRequest(
topic=topic,
message_type=message_type,
filename=filename,
timeout_ms=timeout_ms or int(_DEFAULT_CAPTURE_TIMEOUT_SEC * 1000),
),
timeout,
)
return CaptureImageResult(
filename=response.filename,
file_path=response.file_path,
file_size=int(response.file_size),
mime_type=response.mime_type,
)
def capture_video(
self,
*,
topic: str,
quality: str = "",
filename: str = "",
message_type: str = "",
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> VideoSession:
response = self._call(
self._session.archive_stub.StartVideoCapture,
archive_pb2.StartVideoCaptureRequest(
topic=topic,
quality=quality,
filename=filename,
message_type=message_type,
),
timeout,
)
return VideoSession(self, response.recording_id, response.filename, response.file_path)
def stop_video_capture(
self,
recording_id: str,
*,
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> StopVideoCaptureResult:
response = self._call(
self._session.archive_stub.StopVideoCapture,
archive_pb2.StopVideoCaptureRequest(recording_id=recording_id),
timeout,
)
return StopVideoCaptureResult(
recording_id=response.recording_id,
filename=response.filename,
file_path=response.file_path,
file_size=int(response.file_size),
frame_count=int(response.frame_count),
duration_seconds=float(response.duration_seconds),
reason=response.reason,
)
def list_images(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[ImageSummary]:
response = self._call(
self._session.archive_stub.ListImages,
archive_pb2.ListImagesRequest(),
timeout,
)
return [_image_summary_from_proto(item) for item in response.images]
def list_videos(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[VideoSummary]:
response = self._call(
self._session.archive_stub.ListVideos,
archive_pb2.ListVideosRequest(),
timeout,
)
return [_video_summary_from_proto(item) for item in response.videos]
def list_bags(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[BagSummary]:
response = self._call(
self._session.archive_stub.ListBags,
archive_pb2.ListBagsRequest(),
timeout,
)
return [_bag_summary_from_proto(item) for item in response.bags]
def inspect_bag(self, name: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> BagInfo:
response = self._call(
self._session.archive_stub.InspectBag,
archive_pb2.InspectBagRequest(name=name),
timeout,
)
return _bag_from_proto(response.bag)
def record_bag(
self,
*,
robot_namespace: str = "",
topics: list[str] | None = None,
filename: str = "",
compression: str = "",
storage_format: str = "",
qos_overrides_yaml: str = "",
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> BagSession:
response = self._call(
self._session.archive_stub.StartBagRecording,
archive_pb2.StartBagRecordingRequest(
robot_namespace=normalize_robot_namespace(robot_namespace),
topics=topics or [],
filename=filename,
compression=compression,
storage_format=storage_format,
qos_overrides_yaml=qos_overrides_yaml,
),
timeout,
)
return BagSession(self, response.recording_id, response.filename, response.output_path)
def stop_bag_recording(
self, recording_id: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC
) -> StopBagRecordingResult:
response = self._call(
self._session.archive_stub.StopBagRecording,
archive_pb2.StopBagRecordingRequest(recording_id=recording_id),
timeout,
)
return StopBagRecordingResult(
recording_id=response.recording_id,
filename=response.filename,
total_bytes=int(response.total_bytes),
)
def play_bag(
self,
name: str,
*,
topics: list[str] | None = None,
playback_rate: float = 0.0,
loop: bool = False,
start_offset_seconds: float = 0.0,
clock_type: str = "",
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> PlaybackSession:
response = self._call(
self._session.archive_stub.StartBagPlayback,
archive_pb2.StartBagPlaybackRequest(
name=name,
topics=topics or [],
playback_rate=playback_rate,
loop=loop,
start_offset_seconds=start_offset_seconds,
clock_type=clock_type,
),
timeout,
)
return PlaybackSession(self, response.playback_id, response.name)
def stop_bag_playback(
self, playback_id: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC
) -> StopBagPlaybackResult:
response = self._call(
self._session.archive_stub.StopBagPlayback,
archive_pb2.StopBagPlaybackRequest(playback_id=playback_id),
timeout,
)
return StopBagPlaybackResult(playback_id=response.playback_id)
def pause_bag_playback(self, playback_id: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> None:
self._call(
self._session.archive_stub.PauseBagPlayback,
archive_pb2.PauseBagPlaybackRequest(playback_id=playback_id),
timeout,
)
def resume_bag_playback(self, playback_id: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> None:
self._call(
self._session.archive_stub.ResumeBagPlayback,
archive_pb2.ResumeBagPlaybackRequest(playback_id=playback_id),
timeout,
)
def seek_bag_playback(
self,
playback_id: str,
seek_offset_seconds: float,
*,
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> None:
self._call(
self._session.archive_stub.SeekBagPlayback,
archive_pb2.SeekBagPlaybackRequest(
playback_id=playback_id,
seek_offset_seconds=seek_offset_seconds,
),
timeout,
)
def watch_bag_session(
self, session_id: str, *, timeout: float = _DEFAULT_WATCH_TIMEOUT_SEC
) -> BagSessionTerminalEvent:
stream = self._session.archive_stub.WatchBagSession(
archive_pb2.WatchBagSessionRequest(session_id=session_id),
timeout=timeout,
)
for event in stream:
if event.HasField("started"):
continue
return _terminal_from_event(event)
raise ValueError("Session watch ended without terminal event")
def handle(self, namespace: Optional[str] = None) -> "ArchiveHandle":
return ArchiveHandle(self, namespace)
[docs]
class ArchiveHandle:
"""Namespace-scoped archive handle (symmetry with other robot modules)."""
[docs]
def __init__(self, client: Archive, namespace: Optional[str] = None) -> None:
self._client = client
self._namespace = namespace
def capture_image(
self,
*,
topic: str,
message_type: str,
filename: str = "",
timeout_ms: int = 0,
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> CaptureImageResult:
return self._client.capture_image(
topic=topic,
message_type=message_type,
filename=filename,
timeout_ms=timeout_ms,
timeout=timeout,
)
def capture_video(
self,
*,
topic: str,
quality: str = "",
filename: str = "",
message_type: str = "",
timeout: float = _DEFAULT_TIMEOUT_SEC,
) -> VideoSession:
return self._client.capture_video(
topic=topic,
quality=quality,
filename=filename,
message_type=message_type,
timeout=timeout,
)
def list_images(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[ImageSummary]:
return self._client.list_images(timeout=timeout)
def list_videos(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[VideoSummary]:
return self._client.list_videos(timeout=timeout)
def list_bags(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[BagSummary]:
return self._client.list_bags(timeout=timeout)
def inspect_bag(self, name: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> BagInfo:
return self._client.inspect_bag(name, timeout=timeout)
def record_bag(self, **kwargs) -> BagSession:
if self._namespace is not None:
kwargs.setdefault("robot_namespace", normalize_robot_namespace(self._namespace))
return self._client.record_bag(**kwargs)
def play_bag(self, name: str, **kwargs) -> PlaybackSession:
return self._client.play_bag(name, **kwargs)
__all__ = [
"Archive",
"ArchiveHandle",
"BagInfo",
"BagSession",
"BagSessionTerminalEvent",
"BagSummary",
"BagTopicInfo",
"CaptureImageResult",
"ImageSummary",
"PlaybackSession",
"StopBagPlaybackResult",
"StopBagRecordingResult",
"StopVideoCaptureResult",
"VideoSession",
"VideoSummary",
]