use existing types in CartesianCoords

Author: rickwierenga

pylabrobot/arms/coords.py

@@ -2,6 +2,8 @@
 from enum import Enum
 from typing import Optional
 
+from pylabrobot.resources import Coordinate, Rotation
+
 
 class ElbowOrientation(Enum):
   RIGHT = "right"
@@ -20,10 +22,6 @@ class JointCoords:
 
 @dataclass
 class CartesianCoords:
-  x: float
-  y: float
-  z: float
-  yaw: float
-  pitch: float
-  roll: float
+  location: Coordinate
+  rotation: Rotation
   orientation: Optional[ElbowOrientation] = None

pylabrobot/arms/precise_flex/precise_flex_backend.py

@@ -1,9 +1,12 @@
 from abc import ABC, abstractmethod
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 from pylabrobot.arms.backend import ArmBackend
 from pylabrobot.arms.coords import CartesianCoords, ElbowOrientation, JointCoords
 from pylabrobot.arms.precise_flex.precise_flex_api import PreciseFlexBackendApi
+from pylabrobot.resources import Coordinate, Rotation
+
+PreciseFlexModel = Literal["pf400", "pf3400"]
 
 
 class CoordsConverter(ABC):
@@ -55,7 +58,9 @@ def convert_to_cartesian_space(
       )
     orientation = ElbowOrientation(position[6])
     return CartesianCoords(
-      position[0], position[1], position[2], position[3], position[4], position[5], orientation
+      location=Coordinate(position[0], position[1], position[2]),
+      rotation=Rotation(position[5], position[4], position[3]),
+      orientation=orientation,
     )
 
   def convert_to_joints_array(
@@ -78,12 +83,12 @@ def convert_to_cartesian_array(
     """Convert a CartesianSpace object to a list of cartesian coordinates."""
     orientation_int = self._convert_orientation_enum_to_int(position.orientation)
     arr = (
-      position.x,
-      position.y,
-      position.z,
-      position.yaw,
-      position.pitch,
-      position.roll,
+      position.location.x,
+      position.location.y,
+      position.location.z,
+      position.rotation.yaw,
+      position.rotation.pitch,
+      position.rotation.roll,
       orientation_int,
     )
     return arr
@@ -118,7 +123,9 @@ def convert_to_cartesian_space(
       )
     orientation = ElbowOrientation(position[6])
     return CartesianCoords(
-      position[0], position[1], position[2], position[3], position[4], position[5], orientation
+      location=Coordinate(position[0], position[1], position[2]),
+      rotation=Rotation(position[5], position[4], position[3]),
+      orientation=orientation,
     )
 
   def convert_to_joints_array(
@@ -141,12 +148,12 @@ def convert_to_cartesian_array(
     """Convert a CartesianSpace object to a list of cartesian coordinates."""
     orientation_int = self._convert_orientation_enum_to_int(position.orientation)
     arr = (
-      position.x,
-      position.y,
-      position.z,
-      position.yaw,
-      position.pitch,
-      position.roll,
+      position.location.x,
+      position.location.y,
+      position.location.z,
+      position.rotation.yaw,
+      position.rotation.pitch,
+      position.rotation.roll,
       orientation_int,
     )
     return arr
@@ -164,7 +171,7 @@ def _convert_orientation_enum_to_int(self, orientation: Optional[ElbowOrientatio
 
 class CoordsConverterFactory:
   @staticmethod
-  def create_coords_converter(model: str) -> CoordsConverter:
+  def create_coords_converter(model: PreciseFlexModel) -> CoordsConverter:
     """Factory method to create a CoordsConverter based on the robot model."""
     if model == "pf400":
       return PreciseFlex400SpaceConverter()
@@ -176,7 +183,7 @@ def create_coords_converter(model: str) -> CoordsConverter:
 class PreciseFlexBackend(ArmBackend, ABC):
   """Backend for the PreciseFlex robotic arm  - Default to using Cartesian coordinates, some methods in Brook's TCS don't work with Joint coordinates."""
 
-  def __init__(self, model: str, host: str, port: int = 10100, timeout=20) -> None:
+  def __init__(self, model: PreciseFlexModel, host: str, port: int = 10100, timeout=20) -> None:
     super().__init__()
     self.api = PreciseFlexBackendApi(host=host, port=port, timeout=timeout)
     self.space_converter = CoordsConverterFactory.create_coords_converter(model)
@@ -280,20 +287,20 @@ async def version(self) -> str:
 
   async def approach(self, position: Union[CartesianCoords, JointCoords], approach_height: float):
     """Move the arm to a position above the specified coordinates by a certain distance."""
-    if type(position) == JointCoords:
+    if isinstance(position, JointCoords):
       joints = self.space_converter.convert_to_joints_array(position)
       await self._approach_j(joints, approach_height)
-    elif type(position) == CartesianCoords:
+    elif isinstance(position, CartesianCoords):
       xyz = self.space_converter.convert_to_cartesian_array(position)
       await self._approach_c(xyz[:-1], approach_height, xyz[-1])
     else:
       raise ValueError("Position must be of type JointSpace or CartesianSpace.")
 
   async def pick_plate(self, position: Union[CartesianCoords, JointCoords], approach_height: float):
     """Pick a plate from the specified position."""
-    if type(position) == JointCoords:
+    if isinstance(position, JointCoords):
       raise ValueError("pick_plate only supports CartesianCoords for PreciseFlex.")
-    elif type(position) == CartesianCoords:
+    elif isinstance(position, CartesianCoords):
       xyz = self.space_converter.convert_to_cartesian_array(position)
       await self._pick_plate_c(xyz[:-1], approach_height, xyz[-1])
     else:
@@ -303,20 +310,20 @@ async def place_plate(
     self, position: Union[CartesianCoords, JointCoords], approach_height: float
   ):
     """Place a plate at the specified position."""
-    if type(position) == JointCoords:
+    if isinstance(position, JointCoords):
       raise ValueError("place_plate only supports CartesianCoords for PreciseFlex.")
-    elif type(position) == CartesianCoords:
+    elif isinstance(position, CartesianCoords):
       xyz = self.space_converter.convert_to_cartesian_array(position)
       await self._place_plate_c(xyz[:-1], approach_height, xyz[-1])
     else:
       raise ValueError("Position must be of type JointSpace or CartesianSpace.")
 
   async def move_to(self, position: Union[CartesianCoords, JointCoords]):
     """Move the arm to a specified position in 3D space."""
-    if type(position) == JointCoords:
+    if isinstance(position, JointCoords):
       joints = self.space_converter.convert_to_joints_array(position)
       await self._move_to_j(joints)
-    elif type(position) == CartesianCoords:
+    elif isinstance(position, CartesianCoords):
       xyz = self.space_converter.convert_to_cartesian_array(position)
       await self._move_to_c(xyz[:-1], xyz[-1])
     else:

pylabrobot/arms/precise_flex/precise_flex_backend_tests.py

@@ -4,6 +4,7 @@
 
 from pylabrobot.arms.coords import CartesianCoords, ElbowOrientation, JointCoords
 from pylabrobot.arms.precise_flex.precise_flex_backend import PreciseFlexBackend
+from pylabrobot.resources import Coordinate, Rotation
 
 
 @pytest.mark.hardware  # include/exclude via "pytest -m hardware"
@@ -21,17 +22,25 @@ class PreciseFlexBackendHardwareTests(unittest.IsolatedAsyncioTestCase):
   TEST_PARAMETER_ID = 17018
   TEST_SIGNAL_ID = 20064
 
-  SAFE_LOCATION_C = CartesianCoords(175, 0, 169.994, -0.001, 90, 180, ElbowOrientation.RIGHT)
+  SAFE_LOCATION_C = CartesianCoords(
+    location=Coordinate(175, 0, 169.994),
+    rotation=Rotation(180, 90, -0.001),
+    orientation=ElbowOrientation.RIGHT,
+  )
   SAFE_LOCATION_J = JointCoords(0, 170.003, 0, 180, -180, 75.486)
 
   TEST_LOCATION_J_LEFT = JointCoords(0, 169.932, 16.883, 230.942, -224.288, 75.662)
   TEST_LOCATION_C_LEFT = CartesianCoords(
-    328.426, -115.219, 169.932, 23.537, 90, 180, ElbowOrientation.LEFT
+    location=Coordinate(328.426, -115.219, 169.932),
+    rotation=Rotation(180, 90, 23.537),
+    orientation=ElbowOrientation.LEFT,
   )
 
   TEST_LOCATION_J_RIGHT = JointCoords(0, 169.968, -4.238, 117.915, -100.062, 75.668)
   TEST_LOCATION_C_RIGHT = CartesianCoords(
-    342.562, 280.484, 169.969, 13.612, 90, 180, ElbowOrientation.RIGHT
+    location=Coordinate(342.562, 280.484, 169.969),
+    rotation=Rotation(180, 90, 13.612),
+    orientation=ElbowOrientation.RIGHT,
   )
 
   async def asyncSetUp(self):

pylabrobot/resources/rotation.py

@@ -68,3 +68,15 @@ def deserialize(data) -> "Rotation":
 
   def __repr__(self) -> str:
     return self.__str__()
+
+  @property
+  def roll(self) -> float:
+    return self.x
+
+  @property
+  def pitch(self) -> float:
+    return self.y
+
+  @property
+  def yaw(self) -> float:
+    return self.z