fixed test & linter errors

api/src/opentrons/protocol_engine/execution/labware_movement.py

@@ -1,7 +1,7 @@
 """Labware movement command handling."""
 from __future__ import annotations
 
-from typing import Optional, Union, TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING
 from opentrons_shared_data.gripper.constants import (
     LABWARE_GRIP_FORCE,
     IDLE_STATE_GRIP_FORCE,
@@ -26,11 +26,10 @@
 )
 
 from ..types import (
-    DeckSlotLocation,
-    ModuleLocation,
     OnLabwareLocation,
     LabwareLocation,
     LabwareMovementOffsetData,
+    OnDeckLabwareLocation,
 )
 
 if TYPE_CHECKING:
@@ -84,8 +83,8 @@ def __init__(
     async def move_labware_with_gripper(
         self,
         labware_id: str,
-        current_location: Union[DeckSlotLocation, ModuleLocation, OnLabwareLocation],
-        new_location: Union[DeckSlotLocation, ModuleLocation, OnLabwareLocation],
+        current_location: OnDeckLabwareLocation,
+        new_location: OnDeckLabwareLocation,
         user_offset_data: LabwareMovementOffsetData,
     ) -> None:
         """Move a loaded labware from one location to another using gripper."""

api/src/opentrons/protocol_engine/state/geometry.py

@@ -25,6 +25,7 @@
     CurrentWell,
     TipGeometry,
     LabwareMovementOffsetData,
+    OnDeckLabwareLocation,
 )
 from .config import Config
 from .labware import LabwareView
@@ -36,7 +37,6 @@
 
 
 SLOT_WIDTH = 128
-_ADDITIONAL_TC2_PICKUP_OFFSET = 3.5
 
 
 class _TipDropSection(enum.Enum):
@@ -628,19 +628,19 @@ def _get_drop_tip_well_x_offset(
 
     def get_final_labware_movement_offset_vectors(
         self,
-        from_location: LabwareLocation,
-        to_location: LabwareLocation,
+        from_location: OnDeckLabwareLocation,
+        to_location: OnDeckLabwareLocation,
         additional_offset_vector: LabwareMovementOffsetData,
     ) -> LabwareMovementOffsetData:
         """Calculate the final labware offset vector to use in labware movement."""
         pick_up_offset = (
-            self.get_total_nominal_gripper_offset(
+            self.get_total_nominal_gripper_offset_for_move_type(
                 location=from_location, move_type=_GripperMoveType.PICK_UP_LABWARE
             )
             + additional_offset_vector.pickUpOffset
         )
         drop_offset = (
-            self.get_total_nominal_gripper_offset(
+            self.get_total_nominal_gripper_offset_for_move_type(
                 location=to_location, move_type=_GripperMoveType.DROP_LABWARE
             )
             + additional_offset_vector.dropOffset
@@ -663,65 +663,93 @@ def ensure_valid_gripper_location(
             )
         return location
 
-    def get_total_nominal_gripper_offset(
-        self, location: LabwareLocation, move_type: _GripperMoveType
+    def get_total_nominal_gripper_offset_for_move_type(
+        self, location: OnDeckLabwareLocation, move_type: _GripperMoveType
     ) -> LabwareOffsetVector:
         """Get the total of the offsets to be used to pick up labware in its current location."""
         if move_type == _GripperMoveType.PICK_UP_LABWARE:
             if isinstance(location, (ModuleLocation, DeckSlotLocation)):
-                return self.get_default_gripper_offsets(location).pickUpOffset
-            elif isinstance(location, OnLabwareLocation):
-                # If it's a labware on a labware, we calculate the offset as sum of
-                # offsets for the direct parent labware and the underlying
-                # non-labware parent location.
-                direct_parent_offset = self.get_default_gripper_offsets(location)
+                return self._nominal_gripper_offsets_for_location(location).pickUpOffset
+            else:
+                # If it's a labware on a labware (most likely an adapter),
+                # we calculate the offset as sum of offsets for the direct parent labware
+                # and the underlying non-labware parent location.
+                direct_parent_offset = self._nominal_gripper_offsets_for_location(
+                    location
+                )
                 ancestor = self._labware.get_parent_location(location.labwareId)
+                assert isinstance(
+                    ancestor, (DeckSlotLocation, ModuleLocation)
+                ), "No gripper offsets for off-deck labware"
                 return (
                     direct_parent_offset.pickUpOffset
-                    + self.get_default_gripper_offsets(location=ancestor).pickUpOffset
-                )
-            else:
-                raise errors.LabwareNotOnDeckError(
-                    "No gripper offsets for off-deck labware since Off-deck labware movements"
-                    " are not supported using the gripper."
+                    + self._nominal_gripper_offsets_for_location(
+                        location=ancestor
+                    ).pickUpOffset
                 )
         else:
             if isinstance(location, (ModuleLocation, DeckSlotLocation)):
-                return self.get_default_gripper_offsets(location).dropOffset
-            elif isinstance(location, OnLabwareLocation):
-                # If it's a labware on a labware, we calculate the offset as sum of
-                # offsets for the direct parent labware and the underlying
-                # non-labware parent location.
-                direct_parent_offset = self.get_default_gripper_offsets(location)
+                return self._nominal_gripper_offsets_for_location(location).dropOffset
+            else:
+                # If it's a labware on a labware (most likely an adapter),
+                # we calculate the offset as sum of offsets for the direct parent labware
+                # and the underlying non-labware parent location.
+                direct_parent_offset = self._nominal_gripper_offsets_for_location(
+                    location
+                )
                 ancestor = self._labware.get_parent_location(location.labwareId)
+                assert isinstance(
+                    ancestor, (DeckSlotLocation, ModuleLocation)
+                ), "No gripper offsets for off-deck labware"
                 return (
                     direct_parent_offset.dropOffset
-                    + self.get_default_gripper_offsets(location=ancestor).dropOffset
-                )
-            else:
-                raise errors.LabwareNotOnDeckError(
-                    "No gripper offsets for off-deck labware since Off-deck labware movements"
-                    " are not supported using the gripper."
+                    + self._nominal_gripper_offsets_for_location(
+                        location=ancestor
+                    ).dropOffset
                 )
 
-    def get_default_gripper_offsets(
-        self, location: LabwareLocation
+    def _nominal_gripper_offsets_for_location(
+        self, location: OnDeckLabwareLocation
     ) -> LabwareMovementOffsetData:
         """Provide the default gripper offset data for the given location type."""
         if isinstance(location, DeckSlotLocation):
             offsets = self._labware.get_deck_default_gripper_offsets()
         elif isinstance(location, ModuleLocation):
             offsets = self._modules.get_default_gripper_offsets(location.moduleId)
-        elif isinstance(location, OnLabwareLocation):
-            offsets = self._labware.get_labware_default_gripper_offsets(
-                location.labwareId
-            )
         else:
-            raise errors.LabwareNotOnDeckError(
-                "No gripper offsets for off-deck labware since Off-deck labware movements"
-                " are not supported using the gripper."
-            )
+            # Labware is on a labware/adapter
+            offsets = self._labware_gripper_offsets(location.labwareId)
         return offsets or LabwareMovementOffsetData(
             pickUpOffset=LabwareOffsetVector(x=0, y=0, z=0),
             dropOffset=LabwareOffsetVector(x=0, y=0, z=0),
         )
+
+    def _labware_gripper_offsets(
+        self, labware_id: str
+    ) -> Optional[LabwareMovementOffsetData]:
+        """Provide the most appropriate gripper offset data for the specified labware.
+
+        We check the types of gripper offsets available for the labware ("default" or slot-based)
+        and return the most appropriate one for the overall location of the labware.
+        Currently, only module adapters (specifically, the H/S universal flat adapter)
+        have non-default offsets that are specific to location of the module on deck,
+        so, this code only checks for the presence of those known offsets.
+        """
+        parent_location = self._labware.get_parent_location(labware_id)
+        assert isinstance(
+            parent_location, (DeckSlotLocation, ModuleLocation)
+        ), "No gripper offsets for off-deck labware"
+
+        if isinstance(parent_location, DeckSlotLocation):
+            slot_name = parent_location.slotName
+        else:
+            module_loc = self._modules.get_location(parent_location.moduleId)
+            slot_name = module_loc.slotName
+
+        slot_based_offset = self._labware.get_labware_gripper_offsets(
+            labware_id=labware_id, slot_name=slot_name
+        )
+
+        return slot_based_offset or self._labware.get_labware_gripper_offsets(
+            labware_id=labware_id, slot_name=None
+        )

api/src/opentrons/protocol_engine/state/labware.py

@@ -743,18 +743,21 @@ def get_deck_default_gripper_offsets(self) -> Optional[LabwareMovementOffsetData
             else None
         )
 
-    def get_labware_default_gripper_offsets(
-        self, labware_id: str
+    def get_labware_gripper_offsets(
+        self,
+        labware_id: str,
+        slot_name: Optional[DeckSlotName],
     ) -> Optional[LabwareMovementOffsetData]:
-        """Get the labware's default gripper offsets."""
+        """Get the labware's gripper offsets of the specified type."""
         parsed_offsets = self.get_definition(labware_id).gripperOffsets
+        offset_key = slot_name.name if slot_name else "default"
         return (
             LabwareMovementOffsetData(
                 pickUpOffset=cast(
-                    LabwareOffsetVector, parsed_offsets["default"].pickUpOffset
+                    LabwareOffsetVector, parsed_offsets[offset_key].pickUpOffset
                 ),
                 dropOffset=cast(
-                    LabwareOffsetVector, parsed_offsets["default"].dropOffset
+                    LabwareOffsetVector, parsed_offsets[offset_key].dropOffset
                 ),
             )
             if parsed_offsets

api/src/opentrons/protocol_engine/types.py

@@ -80,6 +80,8 @@ class OnLabwareLocation(BaseModel):
 ]
 """Union of all locations where it's legal to keep a labware."""
 
+OnDeckLabwareLocation = Union[DeckSlotLocation, ModuleLocation, OnLabwareLocation]
+
 NonStackedLocation = Union[DeckSlotLocation, ModuleLocation, _OffDeckLocationType]
 """Union of all locations where it's legal to keep a labware that can't be stacked on another labware"""
 

api/tests/opentrons/protocol_engine/state/test_geometry_view.py

@@ -1488,36 +1488,89 @@ def test_get_total_nominal_gripper_offset(
     )
 
     # Case 1: labware on deck
-    result1 = subject.get_total_nominal_gripper_offset(
+    result1 = subject.get_total_nominal_gripper_offset_for_move_type(
         location=DeckSlotLocation(slotName=DeckSlotName.SLOT_3),
         move_type=_GripperMoveType.PICK_UP_LABWARE,
     )
     assert result1 == LabwareOffsetVector(x=1, y=2, z=3)
 
     # Case 2: labware on module
-    result2 = subject.get_total_nominal_gripper_offset(
+    result2 = subject.get_total_nominal_gripper_offset_for_move_type(
         location=ModuleLocation(moduleId="module-id"),
         move_type=_GripperMoveType.DROP_LABWARE,
     )
     assert result2 == LabwareOffsetVector(x=33, y=22, z=11)
 
 
-def test_get_stacked_labware_total_nominal_offset(
+def test_get_stacked_labware_total_nominal_offset_slot_specific(
     decoy: Decoy,
     labware_view: LabwareView,
     module_view: ModuleView,
     subject: GeometryView,
 ) -> None:
     """Get nominal offset for stacked labware."""
-    # Case: labware on adapter on module
+    # Case: labware on adapter on module, adapter has slot-specific offsets
     decoy.when(module_view.get_default_gripper_offsets("module-id")).then_return(
         LabwareMovementOffsetData(
             pickUpOffset=LabwareOffsetVector(x=11, y=22, z=33),
             dropOffset=LabwareOffsetVector(x=33, y=22, z=11),
         )
     )
+    decoy.when(module_view.get_location("module-id")).then_return(
+        DeckSlotLocation(slotName=DeckSlotName.SLOT_C1)
+    )
+    decoy.when(
+        labware_view.get_labware_gripper_offsets(
+            labware_id="adapter-id", slot_name=DeckSlotName.SLOT_C1
+        )
+    ).then_return(
+        LabwareMovementOffsetData(
+            pickUpOffset=LabwareOffsetVector(x=100, y=200, z=300),
+            dropOffset=LabwareOffsetVector(x=300, y=200, z=100),
+        )
+    )
+    decoy.when(labware_view.get_parent_location("adapter-id")).then_return(
+        ModuleLocation(moduleId="module-id")
+    )
+    result1 = subject.get_total_nominal_gripper_offset_for_move_type(
+        location=OnLabwareLocation(labwareId="adapter-id"),
+        move_type=_GripperMoveType.PICK_UP_LABWARE,
+    )
+    assert result1 == LabwareOffsetVector(x=111, y=222, z=333)
+
+    result2 = subject.get_total_nominal_gripper_offset_for_move_type(
+        location=OnLabwareLocation(labwareId="adapter-id"),
+        move_type=_GripperMoveType.DROP_LABWARE,
+    )
+    assert result2 == LabwareOffsetVector(x=333, y=222, z=111)
+
+
+def test_get_stacked_labware_total_nominal_offset_default(
+    decoy: Decoy,
+    labware_view: LabwareView,
+    module_view: ModuleView,
+    subject: GeometryView,
+) -> None:
+    """Get nominal offset for stacked labware."""
+    # Case: labware on adapter on module, adapter has only default offsets
+    decoy.when(module_view.get_default_gripper_offsets("module-id")).then_return(
+        LabwareMovementOffsetData(
+            pickUpOffset=LabwareOffsetVector(x=11, y=22, z=33),
+            dropOffset=LabwareOffsetVector(x=33, y=22, z=11),
+        )
+    )
+    decoy.when(module_view.get_location("module-id")).then_return(
+        DeckSlotLocation(slotName=DeckSlotName.SLOT_C1)
+    )
+    decoy.when(
+        labware_view.get_labware_gripper_offsets(
+            labware_id="adapter-id", slot_name=DeckSlotName.SLOT_C1
+        )
+    ).then_return(None)
     decoy.when(
-        labware_view.get_labware_default_gripper_offsets("adapter-id")
+        labware_view.get_labware_gripper_offsets(
+            labware_id="adapter-id", slot_name=None
+        )
     ).then_return(
         LabwareMovementOffsetData(
             pickUpOffset=LabwareOffsetVector(x=100, y=200, z=300),
@@ -1527,8 +1580,14 @@ def test_get_stacked_labware_total_nominal_offset(
     decoy.when(labware_view.get_parent_location("adapter-id")).then_return(
         ModuleLocation(moduleId="module-id")
     )
-    result3 = subject.get_total_nominal_gripper_offset(
+    result1 = subject.get_total_nominal_gripper_offset_for_move_type(
         location=OnLabwareLocation(labwareId="adapter-id"),
         move_type=_GripperMoveType.PICK_UP_LABWARE,
     )
-    assert result3 == LabwareOffsetVector(x=111, y=222, z=333)
+    assert result1 == LabwareOffsetVector(x=111, y=222, z=333)
+
+    result2 = subject.get_total_nominal_gripper_offset_for_move_type(
+        location=OnLabwareLocation(labwareId="adapter-id"),
+        move_type=_GripperMoveType.DROP_LABWARE,
+    )
+    assert result2 == LabwareOffsetVector(x=333, y=222, z=111)

api/tests/opentrons/protocol_engine/state/test_labware_view.py

@@ -1305,9 +1305,12 @@ def test_get_labware_gripper_offsets(
         },
     )
 
-    assert subject.get_labware_default_gripper_offsets("plate-id") is None
-    assert subject.get_labware_default_gripper_offsets(
-        "adapter-plate-id"
+    assert (
+        subject.get_labware_gripper_offsets(labware_id="plate-id", slot_name=None)
+        is None
+    )
+    assert subject.get_labware_gripper_offsets(
+        labware_id="adapter-plate-id", slot_name=DeckSlotName.SLOT_D1
     ) == LabwareMovementOffsetData(
         pickUpOffset=LabwareOffsetVector(x=0, y=0, z=0),
         dropOffset=LabwareOffsetVector(x=2, y=0, z=0),