RSDK-4135-2 - add getgeometries protocol (#385)

purplenicole730 · web-flow · commit e2ff53485375 · 2023-08-15T12:31:25.000-04:00
diff --git a/src/viam/components/arm/client.py b/src/viam/components/arm/client.py
@@ -118,4 +118,4 @@ async def get_kinematics(
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/audio_input/client.py b/src/viam/components/audio_input/client.py
@@ -60,4 +60,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/base/client.py b/src/viam/components/base/client.py
@@ -132,4 +132,4 @@ async def do_command(
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/board/client.py b/src/viam/components/board/client.py
@@ -177,4 +177,4 @@ async def set_power_mode(
         await self.client.SetPowerMode(request, timeout=timeout)
 
     async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> List[Geometry]:
-        return await get_geometries(self, extra, timeout)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/camera/client.py b/src/viam/components/camera/client.py
@@ -74,4 +74,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/encoder/client.py b/src/viam/components/encoder/client.py
@@ -73,4 +73,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/gantry/client.py b/src/viam/components/gantry/client.py
@@ -84,4 +84,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/generic/client.py b/src/viam/components/generic/client.py
@@ -33,7 +33,7 @@ async def do_command(self, command: Mapping[str, Any], *, timeout: Optional[floa
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
 
 
 async def do_command(
diff --git a/src/viam/components/gripper/client.py b/src/viam/components/gripper/client.py
@@ -59,4 +59,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/input/client.py b/src/viam/components/input/client.py
@@ -163,4 +163,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/motor/client.py b/src/viam/components/motor/client.py
@@ -144,4 +144,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/movement_sensor/client.py b/src/viam/components/movement_sensor/client.py
@@ -104,4 +104,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/pose_tracker/client.py b/src/viam/components/pose_tracker/client.py
@@ -40,4 +40,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/sensor/client.py b/src/viam/components/sensor/client.py
@@ -34,4 +34,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/components/servo/client.py b/src/viam/components/servo/client.py
@@ -59,4 +59,4 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         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)
+        return await get_geometries(self.client, self.name, extra, timeout)
diff --git a/src/viam/module/types.py b/src/viam/module/types.py
@@ -23,11 +23,3 @@ class Stoppable(Protocol):
 
     def stop(self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None, **kwargs):
         ...
-
-
-@runtime_checkable
-class SupportsGetGeometries(Protocol):
-    """The SupportsGetGeometries protocol defines the requirements for a resource to call get_geometries."""
-
-    def get_geometries(self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None, **kwargs):
-        ...
diff --git a/src/viam/resource/types.py b/src/viam/resource/types.py
@@ -1,6 +1,6 @@
 import re
 import sys
-from typing import TYPE_CHECKING, Callable, ClassVar, Mapping, Sequence
+from typing import TYPE_CHECKING, Callable, ClassVar, Mapping, Optional, Protocol, runtime_checkable, Sequence
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -10,7 +10,7 @@
 from typing_extensions import Self
 
 from viam.proto.app.robot import ComponentConfig
-from viam.proto.common import ResourceName
+from viam.proto.common import ResourceName, GetGeometriesRequest, GetGeometriesResponse
 
 if TYPE_CHECKING:
     from .base import ResourceBase
@@ -198,3 +198,11 @@ def resource_name_from_string(string: str) -> ResourceName:
 
 ResourceCreator: TypeAlias = Callable[[ComponentConfig, Mapping[ResourceName, "ResourceBase"]], "ResourceBase"]
 Validator: TypeAlias = Callable[[ComponentConfig], Sequence[str]]
+
+
+@runtime_checkable
+class SupportsGetGeometries(Protocol):
+    """The SupportsGetGeometries protocol defines the requirements for a resource to call get_geometries."""
+
+    async def GetGeometries(self, request: GetGeometriesRequest, timeout: Optional[float] = None) -> GetGeometriesResponse:
+        ...
diff --git a/src/viam/utils.py b/src/viam/utils.py
@@ -14,7 +14,7 @@
 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.types import Subtype, SupportsGetGeometries
 
 if sys.version_info >= (3, 9):
     from collections.abc import Callable
@@ -159,11 +159,13 @@ def datetime_to_timestamp(dt: datetime) -> Timestamp:
     return timestamp
 
 
-async def get_geometries(client, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> List[Geometry]:
+async def get_geometries(
+    client: SupportsGetGeometries, name: str, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None
+) -> List[Geometry]:
     if extra is None:
         extra = {}
-    request = GetGeometriesRequest(name=client.name, extra=dict_to_struct(extra))
-    response: GetGeometriesResponse = await client.client.GetGeometries(request, timeout=timeout)
+    request = GetGeometriesRequest(name=name, extra=dict_to_struct(extra))
+    response: GetGeometriesResponse = await client.GetGeometries(request, timeout=timeout)
     return [geometry for geometry in response.geometries]
 
 
