viam-sdk 0.3.0__py3-none-linux_armv6l.whl → 0.66.0__py3-none-linux_armv6l.whl

viam/sessions_client.py

@@ -0,0 +1,254 @@
+import asyncio
+import importlib
+import pkgutil
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from copy import deepcopy
+from datetime import timedelta
+from enum import IntEnum
+from threading import Lock, Thread
+from typing import MutableMapping, Optional
+
+from grpclib import Status
+from grpclib.client import Channel
+from grpclib.events import RecvTrailingMetadata, SendRequest, listen
+from grpclib.exceptions import GRPCError, StreamTerminatedError
+from grpclib.metadata import _MetadataLike
+
+from viam import logging
+from viam.gen.common.v1.common_pb2 import safety_heartbeat_monitored
+from viam.proto.robot import RobotServiceStub, SendSessionHeartbeatRequest, StartSessionRequest, StartSessionResponse
+from viam.rpc.dial import DialOptions, dial
+
+LOGGER = logging.getLogger(__name__)
+SESSION_METADATA_KEY = "viam-sid"
+
+
+class _SupportedState(IntEnum):
+    UNKNOWN = 0
+    TRUE = 1
+    FALSE = 2
+
+
+class SessionsClient:
+    """
+    A Session allows a client to express that it is actively connected and
+    supports stopping actuating components when it's not.
+    """
+
+    channel: Channel
+    client: RobotServiceStub
+    _address: str  # direct dial address, when using webRTC this is the local socket rather than a robot address
+    _robot_address: Optional[str]  # the actual machine address on app.viam.com. important for creating a sessions client on Windows
+    _dial_options: DialOptions
+    _disabled: bool
+    _lock: Lock
+    _current_id: str
+    _heartbeat_interval: Optional[timedelta]
+    _supported: _SupportedState
+    _thread: Optional[Thread]
+    _pool: ThreadPoolExecutor
+
+    _HEARTBEAT_MONITORED_METHODS: MutableMapping[str, bool] = {}
+
+    def __init__(
+        self,
+        channel: Channel,
+        direct_dial_address: str,
+        dial_options: Optional[DialOptions],
+        *,
+        disabled: bool = False,
+        robot_addr: Optional[str] = None,
+    ):
+        self.channel = channel
+        self.client = RobotServiceStub(channel)
+        self._address = direct_dial_address
+        self._robot_address = robot_addr
+        self._disabled = disabled
+        self._dial_options = deepcopy(dial_options) if dial_options is not None else DialOptions()
+        if sys.platform != "win32" and sys.platform != "cygwin":
+            self._dial_options.disable_webrtc = True
+        self._lock = Lock()
+        self._current_id = ""
+        self._heartbeat_interval = None
+        self._supported = _SupportedState.UNKNOWN
+        self._thread = None
+        self._pool = ThreadPoolExecutor()
+
+        listen(self.channel, SendRequest, self._send_request)
+        listen(self.channel, RecvTrailingMetadata, self._recv_trailers)
+
+    def reset(self):
+        with self._lock:
+            self._reset()
+
+    def _reset(self):
+        LOGGER.debug("resetting session")
+        self._supported = _SupportedState.UNKNOWN
+        self._current_id = ""
+        self._heartbeat_interval = None
+        if self._thread is not None:
+            try:
+                self._thread.join(timeout=1)
+            except RuntimeError:
+                LOGGER.debug("failed to join session heartbeat thread")
+            self._thread = None
+
+    async def _send_request(self, event: SendRequest):
+        if self._disabled:
+            return
+
+        if not self._is_safety_heartbeat_monitored(event.method_name):
+            return
+
+        event.metadata.update(await self.metadata)
+
+    async def _recv_trailers(self, event: RecvTrailingMetadata):
+        if event.status == Status.INVALID_ARGUMENT and event.status_message == "SESSION_EXPIRED":
+            LOGGER.debug("Session expired")
+            self.reset()
+
+    @asynccontextmanager
+    async def _acquire_lock_async(self):
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(self._pool, self._lock.acquire)
+        try:
+            yield
+        finally:
+            self._lock.release()
+
+    @property
+    async def metadata(self) -> _MetadataLike:
+        async with self._acquire_lock_async():
+            if self._disabled or self._supported != _SupportedState.UNKNOWN:
+                return self._metadata
+
+            request = StartSessionRequest(resume=self._current_id)
+            try:
+                response: StartSessionResponse = await self.client.StartSession(request)
+            except GRPCError as error:
+                if error.status == Status.UNIMPLEMENTED:
+                    self._reset()
+                    self._supported = _SupportedState.FALSE
+                    return self._metadata
+                else:
+                    raise
+
+            if response is None:
+                raise GRPCError(status=Status.INTERNAL, message="Expected response to start session")
+
+            if response.heartbeat_window is None:
+                raise GRPCError(status=Status.INTERNAL, message="Expected heartbeat window in response to start session")
+
+            self._supported = _SupportedState.TRUE
+            self._heartbeat_interval = response.heartbeat_window.ToTimedelta()
+            self._current_id = response.id
+
+            # tick once to ensure heartbeats are supported
+            await self._heartbeat_tick(self.client)
+
+            if self._thread is not None:
+                self._reset()
+            if self._supported == _SupportedState.TRUE:
+                # We send heartbeats faster than the interval window to
+                # ensure that we don't fall outside of it and expire the session.
+                wait = self._heartbeat_interval.total_seconds() / 5
+
+                self._thread = Thread(
+                    name="heartbeat-thread",
+                    target=asyncio.run,
+                    args=(self._heartbeat_process(wait),),
+                    daemon=True,
+                )
+                self._thread.start()
+
+            return self._metadata
+
+    async def _heartbeat_tick(self, client: RobotServiceStub):
+        if not self._current_id:
+            LOGGER.debug("Failed to send heartbeat, session client reset")
+            return
+        request = SendSessionHeartbeatRequest(id=self._current_id)
+
+        try:
+            await client.SendSessionHeartbeat(request)
+        except (GRPCError, StreamTerminatedError):
+            LOGGER.debug("Heartbeat terminated", exc_info=True)
+            self._reset()
+        else:
+            LOGGER.debug("Sent heartbeat successfully")
+
+    def _get_local_addr(self) -> str:
+        if sys.platform != "win32" and sys.platform != "cygwin":
+            # if we're not on windows, we want the direct dial address
+            return self._address
+
+        # return `robot_address` if it exists, otherwise fallback
+        # when using TCP (i.e., on Windows), we need to create a connection to the actual
+        # robot address for a sessions client to maintain connectivity successfully
+        return self._robot_address if self._robot_address is not None else self._address
+
+    async def _heartbeat_process(self, wait: float):
+        addr = self._get_local_addr()
+        channel = await dial(address=addr, options=self._dial_options)
+        client = RobotServiceStub(channel.channel)
+        while True:
+            async with self._acquire_lock_async():
+                if self._supported != _SupportedState.TRUE:
+                    return
+                await self._heartbeat_tick(client)
+            await asyncio.sleep(wait)
+
+    @property
+    def _metadata(self) -> _MetadataLike:
+        if self._supported == _SupportedState.TRUE and self._current_id != "":
+            return {SESSION_METADATA_KEY: self._current_id}
+
+        return {}
+
+    def _is_safety_heartbeat_monitored(self, method: str) -> bool:
+        if method in self._HEARTBEAT_MONITORED_METHODS:
+            return self._HEARTBEAT_MONITORED_METHODS[method]
+
+        parts = method.split("/")
+        if len(parts) != 3:
+            self._HEARTBEAT_MONITORED_METHODS[method] = False
+            return False
+        service_path = parts[1]
+        method_name = parts[2]
+
+        parts = service_path.split(".")
+        if len(parts) < 5:
+            self._HEARTBEAT_MONITORED_METHODS[method] = False
+            return False
+        if parts[0] != "viam":
+            self._HEARTBEAT_MONITORED_METHODS[method] = False
+            return False
+        resource_type = parts[1]
+        resource_subtype = parts[2]
+        version = parts[3]
+        service_name = parts[4]
+        try:
+            module = importlib.import_module(f"viam.gen.{resource_type}.{resource_subtype}.{version}")
+            submods = pkgutil.iter_modules(module.__path__)
+            for mod in submods:
+                if "_pb2" in mod.name:
+                    submod = getattr(module, mod.name)
+                    DESCRIPTOR = getattr(submod, "DESCRIPTOR")
+                    for service in DESCRIPTOR.services_by_name.values():
+                        if service.name == service_name:
+                            for method_actual in service.methods:
+                                if method_actual.name == method_name:
+                                    options = method_actual.GetOptions()
+                                    if options.HasExtension(safety_heartbeat_monitored):
+                                        is_monitored = options.Extensions[safety_heartbeat_monitored]
+                                        self._HEARTBEAT_MONITORED_METHODS[method] = is_monitored
+                                        return is_monitored
+                                    self._HEARTBEAT_MONITORED_METHODS[method] = False
+                                    return False
+            self._HEARTBEAT_MONITORED_METHODS[method] = False
+            return False
+        except Exception:
+            self._HEARTBEAT_MONITORED_METHODS[method] = False
+            return False

viam/streams.py

@@ -0,0 +1,44 @@
+import sys
+
+if sys.version_info >= (3, 9):
+    from collections.abc import AsyncIterator
+else:
+    from typing import AsyncIterator
+
+from typing import Protocol, TypeVar
+
+StreamType = TypeVar("StreamType", covariant=True)
+
+
+class Stream(Protocol[StreamType]):
+    async def next(self) -> StreamType: ...
+
+    def __aiter__(self) -> AsyncIterator:
+        return self
+
+    async def __anext__(self) -> StreamType:
+        return await self.next()
+
+
+class StreamReader(Protocol[StreamType]):
+    async def read(self) -> StreamType: ...
+
+
+class StreamSource(Protocol[StreamType]):
+    async def stream(self) -> Stream[StreamType]: ...
+
+
+class StreamWithIterator(Stream[StreamType]):
+    _stream: AsyncIterator[StreamType]
+
+    def __init__(self, stream: AsyncIterator[StreamType]):
+        self._stream = stream
+
+    async def next(self) -> StreamType:
+        return await self._stream.__anext__()
+
+    def __aiter__(self):
+        return self._stream
+
+    async def __anext__(self) -> StreamType:
+        return await self._stream.__anext__()

viam/utils.py

@@ -3,16 +3,20 @@ import contextvars
 import functools
 import sys
 import threading
-from typing import Any, Dict, List, Mapping, SupportsBytes, SupportsFloat, Type, TypeVar, Union
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Mapping, Optional, SupportsBytes, SupportsFloat, Type, TypeVar, Union
 
 from google.protobuf.json_format import MessageToDict, ParseDict
 from google.protobuf.message import Message
 from google.protobuf.struct_pb2 import ListValue, Struct, Value
+from google.protobuf.timestamp_pb2 import Timestamp
 
-from viam.proto.common import GeoPoint, Orientation, ResourceName, Vector3
+from viam.proto.app.data import CaptureInterval, Filter, TagsFilter
+from viam.proto.common import Geometry, GeoPoint, GetGeometriesRequest, GetGeometriesResponse, Orientation, ResourceName, Vector3
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry
-from viam.resource.types import Subtype
+from viam.resource.rpc_client_base import ResourceRPCClientBase
+from viam.resource.types import API, SupportsGetGeometries
 
 if sys.version_info >= (3, 9):
     from collections.abc import Callable
@@ -28,6 +32,9 @@ else:
 ValueTypes = Union[bool, SupportsBytes, SupportsFloat, List, Mapping, str, None]
 """Types that can be encoded into a protobuf `Value`"""
 
+SensorReading = Union[ValueTypes, Vector3, GeoPoint, Orientation]
+"""Types that can be returned from a sensor"""
+
 
 def primitive_to_value(v: ValueTypes) -> Value:
     """
@@ -97,14 +104,10 @@ def resource_names_for_resource(resource: ResourceBase) -> List[ResourceName]:
     rns: List[ResourceName] = []
 
     for klass in resource.__class__.mro():
-        for registration in Registry.REGISTERED_SUBTYPES().values():
+        for registration in Registry.REGISTERED_APIS().values():
             if klass is registration.resource_type:
-                subtype: Subtype = registration.resource_type.SUBTYPE
-                rns.append(
-                    ResourceName(
-                        namespace=subtype.namespace, type=subtype.resource_type, subtype=subtype.resource_subtype, name=resource.name
-                    )
-                )
+                api: API = registration.resource_type.API
+                rns.append(ResourceName(namespace=api.namespace, type=api.resource_type, subtype=api.resource_subtype, name=resource.name))
     return rns
 
 
@@ -113,8 +116,8 @@ def message_to_struct(message: Message) -> Struct:
     struct.update(
         MessageToDict(
             message,
-            including_default_value_fields=True,
-            preserving_proto_field_name=True,
+            True,
+            True,
         ),
     )
     return struct
@@ -128,7 +131,7 @@ def struct_to_message(struct: Struct, message_type: Type[_T]) -> _T:
     return ParseDict(dct, message_type())
 
 
-def dict_to_struct(obj: Mapping[str, ValueTypes]) -> Struct:
+def dict_to_struct(obj: Optional[Mapping[str, ValueTypes]]) -> Struct:
     def _convert(v: ValueTypes) -> Any:
         if isinstance(v, bool):
             return v
@@ -142,6 +145,8 @@ def dict_to_struct(obj: Mapping[str, ValueTypes]) -> Struct:
             return {k: _convert(vv) for (k, vv) in v.items()}
         return v
 
+    if obj is None:
+        obj = {}
     struct = Struct()
     struct.update({k: _convert(v) for (k, v) in obj.items()})
     return struct
@@ -151,7 +156,28 @@ def struct_to_dict(struct: Struct) -> Dict[str, ValueTypes]:
     return {key: value_to_primitive(value) for (key, value) in struct.fields.items()}
 
 
-def sensor_readings_native_to_value(readings: Mapping[str, Any]) -> Mapping[str, Any]:
+def datetime_to_timestamp(dt: Optional[datetime]) -> Optional[Timestamp]:
+    if dt is None:
+        return None
+    timestamp = Timestamp()
+    timestamp.FromDatetime(dt)
+    return timestamp
+
+
+async def get_geometries(
+    client: SupportsGetGeometries,
+    name: str,
+    extra: Optional[Dict[str, Any]] = None,
+    timeout: Optional[float] = None,
+    metadata: ResourceRPCClientBase.Metadata = ResourceRPCClientBase.Metadata(),
+) -> List[Geometry]:
+    md = metadata.proto
+    request = GetGeometriesRequest(name=name, extra=dict_to_struct(extra))
+    response: GetGeometriesResponse = await client.GetGeometries(request, timeout=timeout, metadata=md)
+    return [geometry for geometry in response.geometries]
+
+
+def sensor_readings_native_to_value(readings: Mapping[str, Any]) -> Mapping[str, Value]:
     prim_readings = dict(readings)
     for key, reading in readings.items():
         if isinstance(reading, Vector3):
@@ -169,7 +195,7 @@ def sensor_readings_native_to_value(readings: Mapping[str, Any]) -> Mapping[str,
     return {key: primitive_to_value(value) for (key, value) in prim_readings.items()}
 
 
-def sensor_readings_value_to_native(readings: Mapping[str, Value]) -> Mapping[str, Any]:
+def sensor_readings_value_to_native(readings: Mapping[str, Value]) -> Mapping[str, SensorReading]:
     prim_readings: Dict[str, Any] = {key: value_to_primitive(value) for (key, value) in readings.items()}
     for key, reading in prim_readings.items():
         if isinstance(reading, Mapping):
@@ -235,3 +261,105 @@ async def to_thread(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs)
     ctx = contextvars.copy_context()
     func_call = functools.partial(ctx.run, func, *args, **kwargs)
     return await loop.run_in_executor(None, func_call)  # type: ignore
+
+
+def from_dm_from_extra(extra: Optional[Dict[str, Any]]) -> bool:
+    """Used in modular filter components to get the 'fromDataManagement' value from an extra map."""
+    if extra is None:
+        return False
+
+    return bool(extra.get("fromDataManagement", False))
+
+
+def create_filter(
+    component_name: Optional[str] = None,
+    component_type: Optional[str] = None,
+    method: Optional[str] = None,
+    robot_name: Optional[str] = None,
+    robot_id: Optional[str] = None,
+    part_name: Optional[str] = None,
+    part_id: Optional[str] = None,
+    location_ids: Optional[List[str]] = None,
+    organization_ids: Optional[List[str]] = None,
+    mime_type: Optional[List[str]] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+    tags: Optional[List[str]] = None,
+    bbox_labels: Optional[List[str]] = None,
+    dataset_id: Optional[str] = None,
+) -> Filter:
+    """Create a `Filter`.
+
+    Args:
+        component_name (Optional[str]): Optional name of the component that captured the data being filtered (for example, "left_motor").
+        component_type (Optional[str]): Optional type of the componenet that captured the data being filtered (for example, "motor").
+        method (Optional[str]): Optional name of the method used to capture the data being filtered (for example, "IsPowered").
+        robot_name (Optional[str]): Optional name of the robot associated with the data being filtered (for example, "viam_rover_1").
+        robot_id (Optional[str]): Optional ID of the robot associated with the data being filtered.
+        part_name (Optional[str]): Optional name of the system part associated with the data being filtered (for example,
+            "viam_rover_1-main").
+        part_id (Optional[str]): Optional ID of the system part associated with the data being filtered.
+        location_ids (Optional[List[str]]): Optional list of location IDs associated with the data being filtered.
+        organization_ids (Optional[List[str]]): Optional list of organization IDs associated with the data being filtered.
+        mime_type (Optional[List[str]]): Optional mime type of data being filtered (for example, "image/png").
+        start_time (Optional[datetime.datetime]): Optional start time of an interval to filter data by.
+        end_time (Optional[datetime.datetime]): Optional end time of an interval to filter data by.
+        tags (Optional[List[str]]): Optional list of tags attached to the data being filtered (for example, ["test"]).
+        bbox_labels (Optional[List[str]]): Optional list of bounding box labels attached to the data being filtered (for example, ["square",
+            "circle"]).
+        dataset_id (Optional[str]): Optional ID of dataset associated with data being filtered
+
+    Returns:
+        viam.proto.app.data.Filter: The `Filter` object.
+    """
+    return Filter(
+        component_name=component_name if component_name else "",
+        component_type=component_type if component_type else "",
+        method=method if method else "",
+        robot_name=robot_name if robot_name else "",
+        robot_id=robot_id if robot_id else "",
+        part_name=part_name if part_name else "",
+        part_id=part_id if part_id else "",
+        location_ids=location_ids,
+        organization_ids=organization_ids,
+        mime_type=mime_type,
+        interval=(
+            CaptureInterval(
+                start=datetime_to_timestamp(start_time),
+                end=datetime_to_timestamp(end_time),
+            )
+        )
+        if start_time or end_time
+        else None,
+        tags_filter=TagsFilter(tags=tags),
+        bbox_labels=bbox_labels,
+        dataset_id=dataset_id if dataset_id else "",
+    )
+
+
+def _alias_param(param_name: str, param_alias: str) -> Callable:
+    """
+    Decorator for aliasing a param in a function. Intended for providing backwards compatibility on params with name changes.
+
+    Args:
+        param_name: name of param in function to alias
+        param_alias: alias that can be used for this param
+    Returns:
+        The input function, plus param alias.
+    """
+
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            alias_param_value = kwargs.get(param_alias)
+            if alias_param_value:
+                # Only use alias value if param is not given.
+                if not kwargs.get(param_name):
+                    kwargs[param_name] = alias_param_value
+                del kwargs[param_alias]
+            result = func(*args, **kwargs)
+            return result
+
+        return wrapper
+
+    return decorator

viam/version_metadata.py

@@ -0,0 +1,4 @@
+__version__ = "0.66.0"
+
+API_VERSION = "v0.1.507"
+SDK_VERSION = __version__