Add missing arm get_geometries method in examples (#401)

Author: cheukt

docs/examples/my_cool_arm.py

@@ -2,10 +2,11 @@
 
 import asyncio
 import json
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from viam.components.arm import Arm, JointPositions, KinematicsFileFormat, Pose
 from viam.operations import run_with_operation
+from viam.proto.common import Capsule, Geometry, Sphere
 
 
 class MyCoolArm(Arm):
@@ -26,6 +27,10 @@ def __init__(self, name: str):
         # Starting joint positions
         self.joint_positions = JointPositions(values=[0, 0, 0, 0, 0, 0])
         self.is_stopped = True
+        self.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)),
+        ]
 
         # Minimal working kinematics model
         self.kinematics = json.dumps(
@@ -92,5 +97,8 @@ async def stop(self, extra: Optional[Dict[str, Any]] = None, **kwargs):
     async def is_moving(self) -> bool:
         return not self.is_stopped
 
+    async def get_geometries(self) -> List[Geometry]:
+        return self.geometries
+
     async def get_kinematics(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[KinematicsFileFormat.ValueType, bytes]:
         return KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, self.kinematics

examples/complex_module/src/arm/my_arm.py

@@ -1,12 +1,13 @@
 import asyncio
 import os
-from typing import Any, ClassVar, Dict, Mapping, Optional, Tuple
+from typing import Any, ClassVar, Dict, List, Mapping, Optional, Tuple
+
 from typing_extensions import Self
 
 from viam.components.arm import Arm, JointPositions, KinematicsFileFormat, Pose
 from viam.operations import run_with_operation
 from viam.proto.app.robot import ComponentConfig
-from viam.proto.common import ResourceName
+from viam.proto.common import Capsule, Geometry, ResourceName, Sphere
 from viam.resource.base import ResourceBase
 from viam.resource.types import Model, ModelFamily
 
@@ -30,6 +31,10 @@ def __init__(self, name: str):
         # Starting joint positions
         self.joint_positions = JointPositions(values=[0, 0, 0, 0, 0, 0])
         self.is_stopped = True
+        self.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)),
+        ]
         super().__init__(name)
 
     @classmethod
@@ -90,6 +95,9 @@ async def stop(self, extra: Optional[Dict[str, Any]] = None, **kwargs):
     async def is_moving(self) -> bool:
         return not self.is_stopped
 
+    async def get_geometries(self) -> List[Geometry]:
+        return self.geometries
+
     async def get_kinematics(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[KinematicsFileFormat.ValueType, bytes]:
         dirname = os.path.dirname(__file__)
         filepath = os.path.join(dirname, "./my_arm_kinematics.json")