refactor(api): Port location-clearing commands to StateUpdate (#16279)

## Overview

Still more incremental work towards EXEC-652.

## Changelog

The following commands clear the current pipette location, sometimes
conditionally. This can be surprising and sneaky, but they need to do
this because they do things like home implicitly.

* `home`
* `retractAxis`
* `thermocycler/openLid`
* `thermocycler/closeLid`
* `heaterShaker/openLabwareLatch`
* `heaterShaker/setAndWaitForShakeSpeed`

They now use the `StateUpdate` mechanism to do this instead of
`PipetteStore` having to mirror their logic. This continues the pattern
started in #16160.

This completes the location update part of EXEC-652. What remains after
this PR is volume updates and pipette loads.

api/src/opentrons/protocol_engine/commands/heater_shaker/open_labware_latch.py

@@ -6,6 +6,7 @@
 
 from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
 from ...errors.error_occurrence import ErrorOccurrence
+from ...state import update_types
 
 if TYPE_CHECKING:
     from opentrons.protocol_engine.state.state import StateView
@@ -54,6 +55,8 @@ async def execute(
         self, params: OpenLabwareLatchParams
     ) -> SuccessData[OpenLabwareLatchResult, None]:
         """Open a Heater-Shaker's labware latch."""
+        state_update = update_types.StateUpdate()
+
         # Allow propagation of ModuleNotLoadedError and WrongModuleTypeError.
         hs_module_substate = self._state_view.modules.get_heater_shaker_module_substate(
             module_id=params.moduleId
@@ -72,6 +75,7 @@ async def execute(
             await self._movement.home(
                 axes=self._state_view.motion.get_robot_mount_axes()
             )
+            state_update.clear_all_pipette_locations()
 
         # Allow propagation of ModuleNotAttachedError.
         hs_hardware_module = self._equipment.get_module_hardware_api(
@@ -84,6 +88,7 @@ async def execute(
         return SuccessData(
             public=OpenLabwareLatchResult(pipetteRetracted=pipette_should_retract),
             private=None,
+            state_update=state_update,
         )
 
 

api/src/opentrons/protocol_engine/commands/heater_shaker/set_and_wait_for_shake_speed.py

@@ -6,6 +6,7 @@
 
 from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
 from ...errors.error_occurrence import ErrorOccurrence
+from ...state import update_types
 
 if TYPE_CHECKING:
     from opentrons.protocol_engine.state.state import StateView
@@ -56,6 +57,8 @@ async def execute(
         params: SetAndWaitForShakeSpeedParams,
     ) -> SuccessData[SetAndWaitForShakeSpeedResult, None]:
         """Set and wait for a Heater-Shaker's target shake speed."""
+        state_update = update_types.StateUpdate()
+
         # Allow propagation of ModuleNotLoadedError and WrongModuleTypeError.
         hs_module_substate = self._state_view.modules.get_heater_shaker_module_substate(
             module_id=params.moduleId
@@ -77,6 +80,7 @@ async def execute(
             await self._movement.home(
                 axes=self._state_view.motion.get_robot_mount_axes()
             )
+            state_update.clear_all_pipette_locations()
 
         # Allow propagation of ModuleNotAttachedError.
         hs_hardware_module = self._equipment.get_module_hardware_api(
@@ -91,6 +95,7 @@ async def execute(
                 pipetteRetracted=pipette_should_retract
             ),
             private=None,
+            state_update=state_update,
         )
 
 

api/src/opentrons/protocol_engine/commands/home.py

@@ -5,6 +5,7 @@
 from typing_extensions import Literal
 
 from opentrons.types import MountType
+from ..state import update_types
 from ..types import MotorAxis
 from .command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
 from ..errors.error_occurrence import ErrorOccurrence
@@ -51,14 +52,21 @@ def __init__(self, movement: MovementHandler, **kwargs: object) -> None:
 
     async def execute(self, params: HomeParams) -> SuccessData[HomeResult, None]:
         """Home some or all motors to establish positional accuracy."""
+        state_update = update_types.StateUpdate()
+
         if (
             params.skipIfMountPositionOk is None
             or not await self._movement.check_for_valid_position(
                 mount=params.skipIfMountPositionOk
             )
         ):
             await self._movement.home(axes=params.axes)
-        return SuccessData(public=HomeResult(), private=None)
+
+        # todo(mm, 2024-09-17): Clearing all pipette locations *unconditionally* is to
+        # preserve prior behavior, but we might only want to do this if we actually home.
+        state_update.clear_all_pipette_locations()
+
+        return SuccessData(public=HomeResult(), private=None, state_update=state_update)
 
 
 class Home(BaseCommand[HomeParams, HomeResult, ErrorOccurrence]):

api/src/opentrons/protocol_engine/commands/move_labware.py

@@ -6,7 +6,9 @@
 from typing_extensions import Literal
 
 from opentrons.types import Point
+from ..state import update_types
 from ..types import (
+    CurrentWell,
     LabwareLocation,
     DeckSlotLocation,
     OnLabwareLocation,
@@ -96,6 +98,8 @@ async def execute(  # noqa: C901
         self, params: MoveLabwareParams
     ) -> SuccessData[MoveLabwareResult, None]:
         """Move a loaded labware to a new location."""
+        state_update = update_types.StateUpdate()
+
         # Allow propagation of LabwareNotLoadedError.
         current_labware = self._state_view.labware.get(labware_id=params.labwareId)
         current_labware_definition = self._state_view.labware.get_definition(
@@ -209,12 +213,28 @@ async def execute(  # noqa: C901
                 user_offset_data=user_offset_data,
                 post_drop_slide_offset=post_drop_slide_offset,
             )
+            # All mounts will have been retracted as part of the gripper move.
+            state_update.clear_all_pipette_locations()
         elif params.strategy == LabwareMovementStrategy.MANUAL_MOVE_WITH_PAUSE:
             # Pause to allow for manual labware movement
             await self._run_control.wait_for_resume()
 
+        # We may have just moved the labware that contains the current well out from
+        # under the pipette. Clear the current location to reflect the fact that the
+        # pipette is no longer over any labware. This is necessary for safe path
+        # planning in case the next movement goes to the same labware (now in a new
+        # place).
+        pipette_location = self._state_view.pipettes.get_current_location()
+        if (
+            isinstance(pipette_location, CurrentWell)
+            and pipette_location.labware_id == params.labwareId
+        ):
+            state_update.clear_all_pipette_locations()
+
         return SuccessData(
-            public=MoveLabwareResult(offsetId=new_offset_id), private=None
+            public=MoveLabwareResult(offsetId=new_offset_id),
+            private=None,
+            state_update=state_update,
         )
 
 

api/src/opentrons/protocol_engine/commands/retract_axis.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Optional, Type
 from typing_extensions import Literal
 
+from ..state import update_types
 from ..types import MotorAxis
 from .command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
 from ..errors.error_occurrence import ErrorOccurrence
@@ -49,8 +50,12 @@ async def execute(
         self, params: RetractAxisParams
     ) -> SuccessData[RetractAxisResult, None]:
         """Retract the specified axis."""
+        state_update = update_types.StateUpdate()
         await self._movement.retract_axis(axis=params.axis)
-        return SuccessData(public=RetractAxisResult(), private=None)
+        state_update.clear_all_pipette_locations()
+        return SuccessData(
+            public=RetractAxisResult(), private=None, state_update=state_update
+        )
 
 
 class RetractAxis(BaseCommand[RetractAxisParams, RetractAxisResult, ErrorOccurrence]):

api/src/opentrons/protocol_engine/commands/thermocycler/close_lid.py

@@ -7,6 +7,7 @@
 
 from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
 from ...errors.error_occurrence import ErrorOccurrence
+from ...state import update_types
 from opentrons.protocol_engine.types import MotorAxis
 
 if TYPE_CHECKING:
@@ -47,6 +48,8 @@ async def execute(
         self, params: CloseLidParams
     ) -> SuccessData[CloseLidResult, None]:
         """Close a Thermocycler's lid."""
+        state_update = update_types.StateUpdate()
+
         thermocycler_state = self._state_view.modules.get_thermocycler_module_substate(
             params.moduleId
         )
@@ -61,11 +64,14 @@ async def execute(
             MotorAxis.Y,
         ] + self._state_view.motion.get_robot_mount_axes()
         await self._movement.home(axes=axes_to_home)
+        state_update.clear_all_pipette_locations()
 
         if thermocycler_hardware is not None:
             await thermocycler_hardware.close()
 
-        return SuccessData(public=CloseLidResult(), private=None)
+        return SuccessData(
+            public=CloseLidResult(), private=None, state_update=state_update
+        )
 
 
 class CloseLid(BaseCommand[CloseLidParams, CloseLidResult, ErrorOccurrence]):

api/src/opentrons/protocol_engine/commands/thermocycler/open_lid.py

@@ -6,6 +6,7 @@
 from pydantic import BaseModel, Field
 
 from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
+from ...state import update_types
 from ...errors.error_occurrence import ErrorOccurrence
 from opentrons.protocol_engine.types import MotorAxis
 
@@ -43,6 +44,8 @@ def __init__(
 
     async def execute(self, params: OpenLidParams) -> SuccessData[OpenLidResult, None]:
         """Open a Thermocycler's lid."""
+        state_update = update_types.StateUpdate()
+
         thermocycler_state = self._state_view.modules.get_thermocycler_module_substate(
             params.moduleId
         )
@@ -57,11 +60,14 @@ async def execute(self, params: OpenLidParams) -> SuccessData[OpenLidResult, Non
             MotorAxis.Y,
         ] + self._state_view.motion.get_robot_mount_axes()
         await self._movement.home(axes=axes_to_home)
+        state_update.clear_all_pipette_locations()
 
         if thermocycler_hardware is not None:
             await thermocycler_hardware.open()
 
-        return SuccessData(public=OpenLidResult(), private=None)
+        return SuccessData(
+            public=OpenLidResult(), private=None, state_update=state_update
+        )
 
 
 class OpenLid(BaseCommand[OpenLidParams, OpenLidResult, ErrorOccurrence]):

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

@@ -152,7 +152,6 @@ def _handle_command(  # noqa: C901
         self, action: Union[SucceedCommandAction, FailCommandAction]
     ) -> None:
         self._update_current_location(action)
-        self._update_deck_point(action)
         self._update_volumes(action)
 
         if not isinstance(action, SucceedCommandAction):
@@ -281,7 +280,7 @@ def _handle_command(  # noqa: C901
                     default_dispense=tip_configuration.default_dispense_flowrate.values_by_api_level,
                 )
 
-    def _update_current_location(  # noqa: C901
+    def _update_current_location(
         self, action: Union[SucceedCommandAction, FailCommandAction]
     ) -> None:
         if isinstance(action, SucceedCommandAction):
@@ -293,8 +292,17 @@ def _update_current_location(  # noqa: C901
             assert_type(action.error, EnumeratedError)
             return
 
-        if location_update != update_types.NO_CHANGE:
-            match location_update.new_location:
+        if location_update is update_types.NO_CHANGE:
+            pass
+        elif location_update is update_types.CLEAR:
+            self._state.current_location = None
+            self._state.current_deck_point = CurrentDeckPoint(
+                mount=None, deck_point=None
+            )
+        else:
+            new_logical_location = location_update.new_location
+            new_deck_point = location_update.new_deck_point
+            match new_logical_location:
                 case update_types.Well(labware_id=labware_id, well_name=well_name):
                     self._state.current_location = CurrentWell(
                         pipette_id=location_update.pipette_id,
@@ -312,110 +320,11 @@ def _update_current_location(  # noqa: C901
                     self._state.current_location = None
                 case update_types.NO_CHANGE:
                     pass
-
-        # todo(mm, 2024-08-29): Port the following isinstance() checks to
-        # use `state_update`. https://opentrons.atlassian.net/browse/EXEC-639
-
-        # These commands leave the pipette in a place that we can't logically associate
-        # with a well. Clear current_location to reflect the fact that it's now unknown.
-        #
-        # TODO(mc, 2021-11-12): Wipe out current_location on movement failures, too.
-        elif isinstance(action, SucceedCommandAction) and isinstance(
-            action.command.result,
-            (
-                commands.HomeResult,
-                commands.RetractAxisResult,
-                commands.thermocycler.OpenLidResult,
-                commands.thermocycler.CloseLidResult,
-            ),
-        ):
-            self._state.current_location = None
-
-        # Heater-Shaker commands may have left the pipette in a place that we can't
-        # associate with a logical location, depending on their result.
-        elif isinstance(action, SucceedCommandAction) and isinstance(
-            action.command.result,
-            (
-                commands.heater_shaker.SetAndWaitForShakeSpeedResult,
-                commands.heater_shaker.OpenLabwareLatchResult,
-            ),
-        ):
-            if action.command.result.pipetteRetracted:
-                self._state.current_location = None
-
-        # A moveLabware command may have moved the labware that contains the current
-        # well out from under the pipette. Clear the current location to reflect the
-        # fact that the pipette is no longer over any labware.
-        #
-        # This is necessary for safe motion planning in case the next movement
-        # goes to the same labware (now in a new place).
-        elif isinstance(action, SucceedCommandAction) and isinstance(
-            action.command.result, commands.MoveLabwareResult
-        ):
-            moved_labware_id = action.command.params.labwareId
-            if action.command.params.strategy == "usingGripper":
-                # All mounts will have been retracted.
-                self._state.current_location = None
-            elif (
-                isinstance(self._state.current_location, CurrentWell)
-                and self._state.current_location.labware_id == moved_labware_id
-            ):
-                self._state.current_location = None
-
-    def _update_deck_point(
-        self, action: Union[SucceedCommandAction, FailCommandAction]
-    ) -> None:
-        if isinstance(action, SucceedCommandAction):
-            location_update = action.state_update.pipette_location
-        elif isinstance(action.error, DefinedErrorData):
-            location_update = action.error.state_update.pipette_location
-        else:
-            # The command failed with some undefined error. We have nothing to do.
-            assert_type(action.error, EnumeratedError)
-            return
-
-        if (
-            location_update is not update_types.NO_CHANGE
-            and location_update.new_deck_point is not update_types.NO_CHANGE
-        ):
-            loaded_pipette = self._state.pipettes_by_id[location_update.pipette_id]
-            self._state.current_deck_point = CurrentDeckPoint(
-                mount=loaded_pipette.mount, deck_point=location_update.new_deck_point
-            )
-
-        # todo(mm, 2024-08-29): Port the following isinstance() checks to
-        # use `state_update`. https://opentrons.atlassian.net/browse/EXEC-639
-        #
-        # These isinstance() checks mostly mirror self._update_current_location().
-        # See there for explanations.
-
-        elif isinstance(action, SucceedCommandAction) and isinstance(
-            action.command.result,
-            (
-                commands.HomeResult,
-                commands.RetractAxisResult,
-                commands.thermocycler.OpenLidResult,
-                commands.thermocycler.CloseLidResult,
-            ),
-        ):
-            self._clear_deck_point()
-
-        elif isinstance(action, SucceedCommandAction) and isinstance(
-            action.command.result,
-            (
-                commands.heater_shaker.SetAndWaitForShakeSpeedResult,
-                commands.heater_shaker.OpenLabwareLatchResult,
-            ),
-        ):
-            if action.command.result.pipetteRetracted:
-                self._clear_deck_point()
-
-        elif isinstance(action, SucceedCommandAction) and isinstance(
-            action.command.result, commands.MoveLabwareResult
-        ):
-            if action.command.params.strategy == "usingGripper":
-                # All mounts will have been retracted.
-                self._clear_deck_point()
+            if new_deck_point is not update_types.NO_CHANGE:
+                loaded_pipette = self._state.pipettes_by_id[location_update.pipette_id]
+                self._state.current_deck_point = CurrentDeckPoint(
+                    mount=loaded_pipette.mount, deck_point=new_deck_point
+                )
 
     def _update_volumes(
         self, action: Union[SucceedCommandAction, FailCommandAction]
@@ -460,10 +369,6 @@ def _update_volumes(
             pipette_id = action.command.params.pipetteId
             self._state.aspirated_volume_by_id[pipette_id] = 0
 
-    def _clear_deck_point(self) -> None:
-        """Reset last deck point to default None value for mount and point."""
-        self._state.current_deck_point = CurrentDeckPoint(mount=None, deck_point=None)
-
 
 class PipetteView(HasState[PipetteState]):
     """Read-only view of computed pipettes state."""