RSDK-4135 - implement getgeometries (#383)

examples/server/v1/components.py

@@ -39,19 +39,27 @@
 from viam.proto.common import (
     AnalogStatus,
     BoardStatus,
+    Capsule,
     DigitalInterruptStatus,
+    Geometry,
     GeoPoint,
     KinematicsFileFormat,
     Orientation,
     Pose,
     PoseInFrame,
-    Vector3,
     ResponseMetadata,
+    Sphere,
+    Vector3,
 )
 from viam.proto.component.arm import JointPositions
 from viam.proto.component.audioinput import AudioChunk, AudioChunkInfo, SampleFormat
 from viam.proto.component.encoder import PositionType
 
+GEOMETRIES = [
+    Geometry(center=Pose(x=1, y=2, z=3, o_x=2, o_y=3, o_z=4, theta=20), sphere=Sphere(radius_mm=2)),
+    Geometry(center=Pose(x=1, y=2, z=3, o_x=2, o_y=3, o_z=4, theta=20), capsule=Capsule(radius_mm=3, length_mm=8)),
+]
+
 
 class ExampleArm(Arm):
     def __init__(self, name: str):
@@ -97,6 +105,9 @@ async def is_moving(self):
     async def get_kinematics(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[KinematicsFileFormat.ValueType, bytes]:
         return self.kinematics
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleAudioInput(AudioInput):
     def __init__(self, name: str):
@@ -160,6 +171,9 @@ async def get_properties(self) -> AudioInput.Properties:
             is_interleaved=True,
         )
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleBase(Base):
     def __init__(self, name: str):
@@ -212,6 +226,9 @@ async def is_moving(self):
     async def get_properties(self, *, timeout: Optional[float] = None, **kwargs) -> Base.Properties:
         return self.props
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleAnalogReader(Board.AnalogReader):
     def __init__(self, name: str, value: int):
@@ -321,6 +338,9 @@ async def model_attributes(self) -> Board.Attributes:
     async def set_power_mode(self, **kwargs):
         raise NotImplementedError()
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleCamera(Camera):
     def __init__(self, name: str):
@@ -343,6 +363,9 @@ async def get_point_cloud(self, **kwargs) -> Tuple[bytes, str]:
     async def get_properties(self, **kwargs) -> Camera.Properties:
         raise NotImplementedError()
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleController(Controller):
     CONTROL_MAP: Dict[int, Control] = {
@@ -448,6 +471,9 @@ def register_control_callback(self, control: Control, triggers: List[EventType],
                     callbacks[trigger] = function
             self.callbacks[control] = callbacks
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleEncoder(Encoder):
     def __init__(self, name: str):
@@ -468,6 +494,9 @@ async def get_position(
     async def get_properties(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Encoder.Properties:
         return Encoder.Properties(ticks_count_supported=True, angle_degrees_supported=False)
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleGantry(Gantry):
     def __init__(self, name: str, position: List[float], lengths: List[float]):
@@ -501,6 +530,9 @@ async def stop(self, extra: Optional[Dict[str, Any]] = None, **kwargs):
     async def is_moving(self):
         return not self.is_stopped
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleGripper(Gripper):
     def __init__(self, name: str):
@@ -523,6 +555,9 @@ async def stop(self, extra: Optional[Dict[str, Any]] = None, **kwargs):
     async def is_moving(self):
         return not self.is_stopped
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleMotor(Motor):
     def __init__(self, name: str):
@@ -596,6 +631,9 @@ async def is_powered(self, extra: Optional[Dict[str, Any]] = None, **kwargs) ->
     async def is_moving(self):
         return self.powered
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleMovementSensor(MovementSensor):
     def __init__(
@@ -646,6 +684,9 @@ async def get_properties(self, **kwargs) -> MovementSensor.Properties:
     async def get_accuracy(self, **kwargs) -> Mapping[str, float]:
         return self.accuracy
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExamplePoseTracker(PoseTracker):
     async def get_poses(self, body_names: List[str], **kwargs) -> Dict[str, PoseInFrame]:
@@ -655,6 +696,9 @@ async def get_poses(self, body_names: List[str], **kwargs) -> Dict[str, PoseInFr
         }
         return {k: v for k, v in all_poses.items() if k in body_names}
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleSensor(Sensor):
     def __init__(self, name: str):
@@ -664,6 +708,9 @@ def __init__(self, name: str):
     async def get_readings(self, **kwargs) -> Mapping[str, Any]:
         return {"abcdefghij"[idx]: random.random() for idx in range(self.num_readings)}
 
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES
+
 
 class ExampleServo(Servo):
     def __init__(self, name: str):
@@ -683,3 +730,6 @@ async def stop(self, **kwargs):
 
     async def is_moving(self):
         return not self.is_stopped
+
+    async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
+        return GEOMETRIES

src/viam/components/arm/__init__.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from viam.proto.common import KinematicsFileFormat, Pose
+from viam.proto.common import Geometry, KinematicsFileFormat, Pose
 from viam.proto.component.arm import JointPositions
 from viam.proto.component.arm import Status as ArmStatus
 from viam.proto.robot import Status
@@ -13,18 +13,15 @@
 
 __all__ = [
     "Arm",
+    "Geometry",
     "JointPositions",
     "KinematicsFileFormat",
     "Pose",
 ]
 
 
 async def create_status(component: Arm) -> Status:
-    (
-        end_position,
-        joint_positions,
-        is_moving,
-    ) = await asyncio.gather(
+    (end_position, joint_positions, is_moving,) = await asyncio.gather(
         component.get_end_position(),
         component.get_joint_positions(),
         component.is_moving(),

src/viam/components/arm/arm.py

@@ -1,10 +1,10 @@
 import abc
-from typing import Any, Dict, Final, Optional, Tuple
+from typing import Any, Dict, Final, List, Optional, Tuple
 
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
 
 from ..component_base import ComponentBase
-from . import JointPositions, KinematicsFileFormat, Pose
+from . import Geometry, JointPositions, KinematicsFileFormat, Pose
 
 
 class Arm(ComponentBase):
@@ -122,3 +122,13 @@ async def get_kinematics(
                 - bytes: The byte contents of the file.
         """
         ...
+
+    @abc.abstractmethod
+    async def get_geometries(self) -> List[Geometry]:
+        """
+        Get all geometries associated with the arm, in their current configuration, in the frame of the arm.
+
+        Returns:
+            List[Geometry]: The geometries associated with the arm.
+        """
+        ...

src/viam/components/arm/client.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Mapping, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Tuple
 
 from grpclib.client import Channel
 
@@ -17,9 +17,9 @@
     StopRequest,
 )
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
-from viam.utils import ValueTypes, dict_to_struct, struct_to_dict
+from viam.utils import ValueTypes, get_geometries, dict_to_struct, struct_to_dict
 
-from . import Arm, KinematicsFileFormat, Pose
+from . import Arm, Geometry, KinematicsFileFormat, Pose
 
 
 class ArmClient(Arm, ReconfigurableResourceRPCClientBase):
@@ -116,3 +116,6 @@ async def get_kinematics(
         request = GetKinematicsRequest(name=self.name, extra=dict_to_struct(extra))
         response: GetKinematicsResponse = await self.client.GetKinematics(request, timeout=timeout)
         return (response.format, response.kinematics_data)
+
+    async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> List[Geometry]:
+        return await get_geometries(self, extra, timeout)

src/viam/components/arm/service.py

@@ -1,6 +1,5 @@
 from grpclib.server import Stream
 
-from viam.errors import MethodNotImplementedError
 from viam.proto.common import (
     DoCommandRequest,
     DoCommandResponse,
@@ -115,4 +114,10 @@ async def GetKinematics(self, stream: Stream[GetKinematicsRequest, GetKinematics
         await stream.send_message(response)
 
     async def GetGeometries(self, stream: Stream[GetGeometriesRequest, GetGeometriesResponse]) -> None:
-        raise MethodNotImplementedError("GetGeometries").grpc_error
+        request = await stream.recv_message()
+        assert request is not None
+        arm = self.get_resource(request.name)
+        timeout = stream.deadline.time_remaining() if stream.deadline else None
+        geometries = await arm.get_geometries(extra=struct_to_dict(request.extra), timeout=timeout)
+        response = GetGeometriesResponse(geometries=geometries)
+        await stream.send_message(response)

src/viam/components/audio_input/__init__.py

@@ -1,3 +1,4 @@
+from viam.proto.common import Geometry
 from viam.resource.registry import Registry, ResourceRegistration
 
 from .audio_input import AudioInput
@@ -6,6 +7,7 @@
 
 __all__ = [
     "AudioInput",
+    "Geometry",
 ]
 
 

src/viam/components/audio_input/audio_input.py

@@ -1,7 +1,7 @@
 import abc
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Final, Optional
+from typing import Final, List, Optional
 
 from google.protobuf.duration_pb2 import Duration
 from typing_extensions import Self
@@ -12,6 +12,7 @@
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
 
 from ..component_base import ComponentBase
+from . import Geometry
 
 
 class AudioInput(ComponentBase, MediaSource[Audio]):
@@ -77,3 +78,13 @@ async def get_properties(self, *, timeout: Optional[float] = None, **kwargs) ->
             Properties: The audio input properties
         """
         ...
+
+    @abc.abstractmethod
+    async def get_geometries(self) -> List[Geometry]:
+        """
+        Get all geometries associated with the audio input, in their current configuration, in the frame of the audio input.
+
+        Returns:
+            List[Geometry]: The geometries associated with the audio input.
+        """
+        ...

src/viam/components/audio_input/client.py

@@ -1,4 +1,4 @@
-from typing import AsyncIterator, Mapping, Optional, Union
+from typing import Any, AsyncIterator, Dict, List, Mapping, Optional, Union
 
 from grpclib.client import Channel
 
@@ -14,9 +14,9 @@
     SampleFormat,
 )
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
-from viam.utils import ValueTypes, dict_to_struct, struct_to_dict
+from viam.utils import ValueTypes, get_geometries, dict_to_struct, struct_to_dict
 
-from .audio_input import AudioInput
+from .audio_input import AudioInput, Geometry
 
 
 class AudioInputClient(AudioInput, ReconfigurableResourceRPCClientBase):
@@ -58,3 +58,6 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         request = DoCommandRequest(name=self.name, command=dict_to_struct(command))
         response: DoCommandResponse = await self.client.DoCommand(request, timeout=timeout)
         return struct_to_dict(response.result)
+
+    async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> List[Geometry]:
+        return await get_geometries(self, extra, timeout)

src/viam/components/audio_input/service.py

@@ -6,7 +6,7 @@
 from grpclib import GRPCError, Status
 from grpclib.server import Stream
 
-from viam.errors import MethodNotImplementedError, NotSupportedError
+from viam.errors import NotSupportedError
 from viam.proto.common import DoCommandRequest, DoCommandResponse, GetGeometriesRequest, GetGeometriesResponse
 from viam.proto.component.audioinput import (
     AudioInputServiceBase,
@@ -105,4 +105,10 @@ async def DoCommand(self, stream: Stream[DoCommandRequest, DoCommandResponse]) -
         await stream.send_message(response)
 
     async def GetGeometries(self, stream: Stream[GetGeometriesRequest, GetGeometriesResponse]) -> None:
-        raise MethodNotImplementedError("GetGeometries").grpc_error
+        request = await stream.recv_message()
+        assert request is not None
+        audio_input = self.get_resource(request.name)
+        timeout = stream.deadline.time_remaining() if stream.deadline else None
+        geometries = await audio_input.get_geometries(extra=struct_to_dict(request.extra), timeout=timeout)
+        response = GetGeometriesResponse(geometries=geometries)
+        await stream.send_message(response)

src/viam/components/base/__init__.py

@@ -1,4 +1,4 @@
-from viam.proto.common import ActuatorStatus, Vector3
+from viam.proto.common import ActuatorStatus, Geometry, Vector3
 from viam.proto.robot import Status
 from viam.resource.registry import Registry, ResourceRegistration
 from viam.utils import message_to_struct
@@ -7,7 +7,11 @@
 from .client import BaseClient
 from .service import BaseRPCService
 
-__all__ = ["Base", "Vector3"]
+__all__ = [
+    "Base",
+    "Geometry",
+    "Vector3",
+]
 
 
 async def create_status(component: Base) -> Status: