[WORKFLOW] AI update based on proto changes from commit bd635a9

Author: viambot

src/viam/components/gantry/client.py

@@ -1,8 +1,8 @@
-from typing import Any, Dict, List, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional, Tuple
 
 from grpclib.client import Channel
 
-from viam.proto.common import DoCommandRequest, DoCommandResponse, Geometry
+from viam.proto.common import DoCommandRequest, DoCommandResponse, Geometry, GetKinematicsRequest, GetKinematicsResponse
 from viam.proto.component.gantry import (
     GantryServiceStub,
     GetLengthsRequest,
@@ -20,6 +20,7 @@
 from viam.utils import ValueTypes, dict_to_struct, get_geometries, struct_to_dict
 
 from .gantry import Gantry
+from .arm import KinematicsFileFormat
 
 
 class GantryClient(Gantry, ReconfigurableResourceRPCClientBase):
@@ -110,6 +111,14 @@ async def do_command(
         response: DoCommandResponse = await self.client.DoCommand(request, timeout=timeout, metadata=md)
         return struct_to_dict(response.result)
 
+    async def get_kinematics(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Tuple[KinematicsFileFormat.ValueType, bytes]:
+        md = kwargs.get("metadata", self.Metadata()).proto
+        request = GetKinematicsRequest(name=self.name, extra=dict_to_struct(extra))
+        response: GetKinematicsResponse = await self.client.GetKinematics(request, timeout=timeout, metadata=md)
+        return response.format, response.kinematics_data
+
     async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs) -> List[Geometry]:
         md = kwargs.get("metadata", self.Metadata())
         return await get_geometries(self.client, self.name, extra, timeout, md)

src/viam/components/gantry/gantry.py

@@ -1,6 +1,7 @@
 import abc
-from typing import Any, Dict, Final, List, Optional
+from typing import Any, Dict, Final, List, Optional, Tuple
 
+from viam.components.arm import KinematicsFileFormat
 from viam.resource.types import API, RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT
 
 from ..component_base import ComponentBase
@@ -154,3 +155,24 @@ async def is_moving(self) -> bool:
         For more information, see `Gantry component <https://docs.viam.com/dev/reference/apis/components/gantry/#ismoving>`_.
         """
         ...
+
+    @abc.abstractmethod
+    async def get_kinematics(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Tuple[KinematicsFileFormat.ValueType, bytes]:
+        """
+        Get the kinematics of the component.
+
+        ::
+
+            my_gantry = Gantry.from_robot(robot=machine, name="my_gantry")
+
+            # Get the kinematics of the gantry.
+            (kinematics_format, kinematics_data) = await my_gantry.get_kinematics()
+
+        Returns:
+            Tuple[KinematicsFileFormat.ValueType, bytes]: The format of the kinematics data and the kinematics data itself.
+
+        For more information, see `Gantry component <https://docs.viam.com/dev/reference/apis/components/gantry/>`_.
+        """
+        ...

src/viam/components/gantry/service.py

@@ -1,6 +1,6 @@
 from grpclib.server import Stream
 
-from viam.proto.common import DoCommandRequest, DoCommandResponse, GetGeometriesRequest, GetGeometriesResponse
+from viam.proto.common import DoCommandRequest, DoCommandResponse, GetGeometriesRequest, GetGeometriesResponse, GetKinematicsRequest, GetKinematicsResponse
 from viam.proto.component.gantry import (
     GantryServiceBase,
     GetLengthsRequest,
@@ -103,6 +103,18 @@ async def DoCommand(self, stream: Stream[DoCommandRequest, DoCommandResponse]) -
         response = DoCommandResponse(result=dict_to_struct(result))
         await stream.send_message(response)
 
+    async def GetKinematics(self, stream: Stream[GetKinematicsRequest, GetKinematicsResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        name = request.name
+        gantry = self.get_resource(name)
+        timeout = stream.deadline.time_remaining() if stream.deadline else None
+        kinematics_format, kinematics_data = await gantry.get_kinematics(
+            extra=struct_to_dict(request.extra), timeout=timeout, metadata=stream.metadata
+        )
+        response = GetKinematicsResponse(format=kinematics_format, kinematics_data=kinematics_data)
+        await stream.send_message(response)
+
     async def GetGeometries(self, stream: Stream[GetGeometriesRequest, GetGeometriesResponse]) -> None:
         request = await stream.recv_message()
         assert request is not None

tests/mocks/components.py

@@ -556,6 +556,7 @@ def __init__(self, name: str, position: List[float], lengths: List[float]):
         self.position = position
         self.lengths = lengths
         self.is_stopped = True
+        self.kinematics = (KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, b"\x00\x01\x02")
         self.extra = None
         self.homed = True
         self.speeds = Optional[List[float]]
@@ -602,6 +603,13 @@ async def stop(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optiona
     async def is_moving(self) -> bool:
         return not self.is_stopped
 
+    async def get_kinematics(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Tuple[KinematicsFileFormat.ValueType, bytes]:
+        self.extra = extra
+        self.timeout = timeout
+        return self.kinematics
+
     async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> List[Geometry]:
         self.extra = extra
         self.timeout = timeout

tests/test_gantry.py

@@ -1,8 +1,9 @@
 from grpclib.testing import ChannelFor
 
+from viam.components.arm import KinematicsFileFormat
 from viam.components.gantry import GantryClient
 from viam.components.gantry.service import GantryRPCService
-from viam.proto.common import DoCommandRequest, DoCommandResponse, GetGeometriesRequest, GetGeometriesResponse
+from viam.proto.common import DoCommandRequest, DoCommandResponse, GetGeometriesRequest, GetGeometriesResponse, GetKinematicsRequest, GetKinematicsResponse
 from viam.proto.component.gantry import (
     GantryServiceStub,
     GetLengthsRequest,
@@ -58,6 +59,11 @@ async def test_is_moving(self):
         await self.gantry.stop()
         assert not await self.gantry.is_moving()
 
+    async def test_get_kinematics(self):
+        kinematics_format, kinematics_data = await self.gantry.get_kinematics()
+        assert kinematics_format == KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA
+        assert kinematics_data == b"\x00\x01\x02"
+
     async def test_extra(self):
         assert self.gantry.extra is None or len(self.gantry.extra) == 0
         extra = {"foo": "bar", "baz": [1, 2, 3]}
@@ -79,6 +85,9 @@ async def test_timeout(self):
         await self.gantry.stop(timeout=4.4)
         assert self.gantry.timeout == loose_approx(4.4)
 
+        await self.gantry.get_kinematics(timeout=9.9)
+        assert self.gantry.timeout == loose_approx(9.9)
+
     async def test_get_geometries(self):
         geometries = await self.gantry.get_geometries()
         assert geometries == GEOMETRIES
@@ -142,6 +151,15 @@ async def test_is_moving(self):
             response: IsMovingResponse = await client.IsMoving(request)
             assert response.is_moving is True
 
+    async def test_get_kinematics(self):
+        async with ChannelFor([self.service]) as channel:
+            client = GantryServiceStub(channel)
+            request = GetKinematicsRequest(name=self.gantry.name)
+            response: GetKinematicsResponse = await client.GetKinematics(request, timeout=1.1)
+            assert response.format == KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA
+            assert response.kinematics_data == b"\x00\x01\x02"
+            assert self.gantry.timeout == loose_approx(1.1)
+
     async def test_extra(self):
         async with ChannelFor([self.service]) as channel:
             assert self.gantry.extra is None or len(self.gantry.extra) == 0
@@ -151,6 +169,11 @@ async def test_extra(self):
             await client.Stop(request)
             assert self.gantry.extra == extra
 
+            extra = {"foo": "bar", "baz": [1, 2, 3]}
+            request = GetKinematicsRequest(name=self.gantry.name, extra=dict_to_struct(extra))
+            await client.GetKinematics(request)
+            assert self.gantry.extra == extra
+
     async def test_do(self):
         async with ChannelFor([self.service]) as channel:
             client = GantryServiceStub(channel)
@@ -226,6 +249,14 @@ async def test_is_moving(self):
             client = GantryClient(self.gantry.name, channel)
             assert await client.is_moving() is True
 
+    async def test_get_kinematics(self):
+        async with ChannelFor([self.service]) as channel:
+            client = GantryClient(self.gantry.name, channel)
+            kinematics_format, kinematics_data = await client.get_kinematics(timeout=2.2)
+            assert kinematics_format == KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA
+            assert kinematics_data == b"\x00\x01\x02"
+            assert self.gantry.timeout == loose_approx(2.2)
+
     async def test_extra(self):
         async with ChannelFor([self.service]) as channel:
             assert self.gantry.extra is None or len(self.gantry.extra) == 0
@@ -234,6 +265,10 @@ async def test_extra(self):
             await client.move_to_position([1, 2, 3], [4, 5, 6], extra=extra)
             assert self.gantry.extra == extra
 
+            extra = {"foo": "bar", "baz": [1, 2, 3]}
+            await client.get_kinematics(extra=extra)
+            assert self.gantry.extra == extra
+
     async def test_get_geometries(self):
         async with ChannelFor([self.service]) as channel:
             client = GantryClient(self.gantry.name, channel)