refactor(api): Port AddressableAreaStore to StateUpdate

api/src/opentrons/protocol_engine/actions/actions.py

@@ -27,7 +27,6 @@
     ModuleDefinition,
     Liquid,
     DeckConfigurationType,
-    AddressableAreaLocation,
 )
 
 
@@ -235,12 +234,12 @@ class SetDeckConfigurationAction:
 class AddAddressableAreaAction:
     """Add a single addressable area to state.
 
-    This differs from the deck configuration in ProvideDeckConfigurationAction which
+    This differs from the deck configuration in SetDeckConfigurationAction which
     sends over a mapping of cutout fixtures. This action will only load one addressable
     area and that should be pre-validated before being sent via the action.
     """
 
-    addressable_area: AddressableAreaLocation
+    addressable_area_name: str
 
 
 @dataclasses.dataclass(frozen=True)

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

@@ -95,7 +95,6 @@ async def execute(self, params: CloseLidParams) -> SuccessData[CloseLidResult]:
                     deck_slot=self._state_view.modules.get_location(
                         params.moduleId
                     ).slotName,
-                    deck_type=self._state_view.config.deck_type,
                     model=absorbance_model,
                 )
             )

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

@@ -91,7 +91,6 @@ async def execute(self, params: OpenLidParams) -> SuccessData[OpenLidResult]:
                     deck_slot=self._state_view.modules.get_location(
                         params.moduleId
                     ).slotName,
-                    deck_type=self._state_view.config.deck_type,
                     model=absorbance_model,
                 )
             )

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

@@ -104,6 +104,8 @@ async def execute(
         self, params: LoadLabwareParams
     ) -> SuccessData[LoadLabwareResult]:
         """Load definition and calibration data necessary for a labware."""
+        state_update = StateUpdate()
+
         # TODO (tz, 8-15-2023): extend column validation to column 1 when working
         # on https://opentrons.atlassian.net/browse/RSS-258 and completing
         # https://opentrons.atlassian.net/browse/RSS-255
@@ -128,10 +130,12 @@ async def execute(
             self._state_view.addressable_areas.raise_if_area_not_in_deck_configuration(
                 area_name
             )
+            state_update.set_addressable_area_used(area_name)
         elif isinstance(params.location, DeckSlotLocation):
             self._state_view.addressable_areas.raise_if_area_not_in_deck_configuration(
                 params.location.slotName.id
             )
+            state_update.set_addressable_area_used(params.location.slotName.id)
 
         verified_location = self._state_view.geometry.ensure_location_not_occupied(
             params.location
@@ -144,8 +148,6 @@ async def execute(
             labware_id=params.labwareId,
         )
 
-        state_update = StateUpdate()
-
         state_update.set_loaded_labware(
             labware_id=loaded_labware.labware_id,
             offset_id=loaded_labware.offsetId,

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

@@ -4,6 +4,8 @@
 from typing_extensions import Literal
 from pydantic import BaseModel, Field
 
+from opentrons.protocol_engine.state.update_types import StateUpdate
+
 from .command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
 from ..errors.error_occurrence import ErrorOccurrence
 from ..types import (
@@ -116,26 +118,35 @@ def __init__(
 
     async def execute(self, params: LoadModuleParams) -> SuccessData[LoadModuleResult]:
         """Check that the requested module is attached and assign its identifier."""
+        state_update = StateUpdate()
+
         module_type = params.model.as_type()
         self._ensure_module_location(params.location.slotName, module_type)
 
+        # todo(mm, 2024-12-03): Theoretically, we should be able to deal with
+        # addressable areas and deck configurations the same way between OT-2 and Flex.
+        # Can this be simplified?
         if self._state_view.config.robot_type == "OT-2 Standard":
             self._state_view.addressable_areas.raise_if_area_not_in_deck_configuration(
                 params.location.slotName.id
             )
         else:
-            addressable_area = self._state_view.geometry._modules.ensure_and_convert_module_fixture_location(
-                deck_slot=params.location.slotName,
-                deck_type=self._state_view.config.deck_type,
-                model=params.model,
+            addressable_area_provided_by_module = (
+                self._state_view.modules.ensure_and_convert_module_fixture_location(
+                    deck_slot=params.location.slotName,
+                    model=params.model,
+                )
             )
             self._state_view.addressable_areas.raise_if_area_not_in_deck_configuration(
-                addressable_area
+                addressable_area_provided_by_module
             )
 
         verified_location = self._state_view.geometry.ensure_location_not_occupied(
             params.location
         )
+        state_update.set_addressable_area_used(
+            addressable_area_name=params.location.slotName.id
+        )
 
         if params.model == ModuleModel.MAGNETIC_BLOCK_V1:
             loaded_module = await self._equipment.load_magnetic_block(
@@ -157,11 +168,15 @@ async def execute(self, params: LoadModuleParams) -> SuccessData[LoadModuleResul
                 model=loaded_module.definition.model,
                 definition=loaded_module.definition,
             ),
+            state_update=state_update,
         )
 
     def _ensure_module_location(
         self, slot: DeckSlotName, module_type: ModuleType
     ) -> None:
+        # todo(mm, 2024-12-03): Theoretically, we should be able to deal with
+        # addressable areas and deck configurations the same way between OT-2 and Flex.
+        # Can this be simplified?
         if self._state_view.config.robot_type == "OT-2 Standard":
             slot_def = self._state_view.addressable_areas.get_slot_definition(slot.id)
             compatible_modules = slot_def["compatibleModuleTypes"]

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

@@ -156,6 +156,7 @@ async def execute(self, params: MoveLabwareParams) -> _ExecuteReturn:  # noqa: C
             self._state_view.addressable_areas.raise_if_area_not_in_deck_configuration(
                 area_name
             )
+            state_update.set_addressable_area_used(addressable_area_name=area_name)
 
             if fixture_validation.is_gripper_waste_chute(area_name):
                 # When dropping off labware in the waste chute, some bigger pieces
@@ -201,6 +202,9 @@ async def execute(self, params: MoveLabwareParams) -> _ExecuteReturn:  # noqa: C
             self._state_view.addressable_areas.raise_if_area_not_in_deck_configuration(
                 params.newLocation.slotName.id
             )
+            state_update.set_addressable_area_used(
+                addressable_area_name=params.newLocation.slotName.id
+            )
 
         available_new_location = self._state_view.geometry.ensure_location_not_occupied(
             location=params.newLocation

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

@@ -265,17 +265,21 @@ async def move_to_addressable_area(
                     )
                 ],
             ),
-            state_update=StateUpdate().clear_all_pipette_locations(),
+            state_update=StateUpdate()
+            .clear_all_pipette_locations()
+            .set_addressable_area_used(addressable_area_name=addressable_area_name),
         )
     else:
         deck_point = DeckPoint.construct(x=x, y=y, z=z)
         return SuccessData(
             public=DestinationPositionResult(position=deck_point),
-            state_update=StateUpdate().set_pipette_location(
+            state_update=StateUpdate()
+            .set_pipette_location(
                 pipette_id=pipette_id,
                 new_addressable_area_name=addressable_area_name,
                 new_deck_point=deck_point,
-            ),
+            )
+            .set_addressable_area_used(addressable_area_name=addressable_area_name),
         )
 
 

api/src/opentrons/protocol_engine/protocol_engine.py

@@ -30,7 +30,6 @@
     HexColor,
     PostRunHardwareState,
     DeckConfigurationType,
-    AddressableAreaLocation,
 )
 from .execution import (
     QueueWorker,
@@ -574,9 +573,8 @@ def add_liquid(
 
     def add_addressable_area(self, addressable_area_name: str) -> None:
         """Add an addressable area to state."""
-        area = AddressableAreaLocation(addressableAreaName=addressable_area_name)
         self._action_dispatcher.dispatch(
-            AddAddressableAreaAction(addressable_area=area)
+            AddAddressableAreaAction(addressable_area_name)
         )
 
     def reset_tips(self, labware_id: str) -> None:

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

@@ -1,7 +1,7 @@
 """Basic addressable area data state and store."""
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Dict, List, Optional, Set, Union
+from typing import Dict, List, Optional, Set
 
 from opentrons_shared_data.robot.types import RobotType, RobotDefinition
 from opentrons_shared_data.deck.types import (
@@ -12,14 +12,6 @@
 
 from opentrons.types import Point, DeckSlotName
 
-from ..commands import (
-    Command,
-    LoadLabwareResult,
-    LoadModuleResult,
-    MoveLabwareResult,
-    MoveToAddressableAreaResult,
-    MoveToAddressableAreaForDropTipResult,
-)
 from ..errors import (
     IncompatibleAddressableAreaError,
     AreaNotInDeckConfigurationError,
@@ -29,19 +21,18 @@
 )
 from ..resources import deck_configuration_provider
 from ..types import (
-    DeckSlotLocation,
-    AddressableAreaLocation,
     AddressableArea,
     PotentialCutoutFixture,
     DeckConfigurationType,
     Dimensions,
 )
+from ..actions.get_state_update import get_state_updates
 from ..actions import (
     Action,
-    SucceedCommandAction,
     SetDeckConfigurationAction,
     AddAddressableAreaAction,
 )
+from . import update_types
 from .config import Config
 from ._abstract_store import HasState, HandlesActions
 
@@ -191,12 +182,18 @@ def __init__(
             robot_definition=robot_definition,
         )
 
+    # TODO: Port loadLabware, moveLabware, loadModule, moveToAddressableArea, and moveToAddressableAreaForDropTip
+    # and their tests
     def handle_action(self, action: Action) -> None:
         """Modify state in reaction to an action."""
-        if isinstance(action, SucceedCommandAction):
-            self._handle_command(action.command)
-        elif isinstance(action, AddAddressableAreaAction):
-            self._check_location_is_addressable_area(action.addressable_area)
+        for state_update in get_state_updates(action):
+            if state_update.addressable_area_used != update_types.NO_CHANGE:
+                self._add_addressable_area(
+                    state_update.addressable_area_used.addressable_area_name
+                )
+
+        if isinstance(action, AddAddressableAreaAction):
+            self._add_addressable_area(action.addressable_area_name)
         elif isinstance(action, SetDeckConfigurationAction):
             current_state = self._state
             if (
@@ -211,28 +208,6 @@ def handle_action(self, action: Action) -> None:
                     )
                 )
 
-    def _handle_command(self, command: Command) -> None:
-        """Modify state in reaction to a command."""
-        if isinstance(command.result, LoadLabwareResult):
-            location = command.params.location
-            if isinstance(location, (DeckSlotLocation, AddressableAreaLocation)):
-                self._check_location_is_addressable_area(location)
-
-        elif isinstance(command.result, MoveLabwareResult):
-            location = command.params.newLocation
-            if isinstance(location, (DeckSlotLocation, AddressableAreaLocation)):
-                self._check_location_is_addressable_area(location)
-
-        elif isinstance(command.result, LoadModuleResult):
-            self._check_location_is_addressable_area(command.params.location)
-
-        elif isinstance(
-            command.result,
-            (MoveToAddressableAreaResult, MoveToAddressableAreaForDropTipResult),
-        ):
-            addressable_area_name = command.params.addressableAreaName
-            self._check_location_is_addressable_area(addressable_area_name)
-
     @staticmethod
     def _get_addressable_areas_from_deck_configuration(
         deck_config: DeckConfigurationType, deck_definition: DeckDefinitionV5
@@ -260,16 +235,7 @@ def _get_addressable_areas_from_deck_configuration(
                 )
         return {area.area_name: area for area in addressable_areas}
 
-    def _check_location_is_addressable_area(
-        self, location: Union[DeckSlotLocation, AddressableAreaLocation, str]
-    ) -> None:
-        if isinstance(location, DeckSlotLocation):
-            addressable_area_name = location.slotName.id
-        elif isinstance(location, AddressableAreaLocation):
-            addressable_area_name = location.addressableAreaName
-        else:
-            addressable_area_name = location
-
+    def _add_addressable_area(self, addressable_area_name: str) -> None:
         if addressable_area_name not in self._state.loaded_addressable_areas_by_name:
             cutout_id = self._validate_addressable_area_for_simulation(
                 addressable_area_name

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

@@ -908,7 +908,7 @@ def get_nominal_offset_to_child(
                     "Module location invalid for nominal module offset calculation."
                 )
             module_addressable_area = self.ensure_and_convert_module_fixture_location(
-                location, self._state.deck_type, module.model
+                location, module.model
             )
             module_addressable_area_position = (
                 addressable_areas.get_addressable_area_offsets_from_cutout(
@@ -1281,13 +1281,14 @@ def convert_absorbance_reader_data_points(
     def ensure_and_convert_module_fixture_location(
         self,
         deck_slot: DeckSlotName,
-        deck_type: DeckType,
         model: ModuleModel,
     ) -> str:
         """Ensure module fixture load location is valid.
 
         Also, convert the deck slot to a valid module fixture addressable area.
         """
+        deck_type = self._state.deck_type
+
         if deck_type == DeckType.OT2_STANDARD or deck_type == DeckType.OT2_SHORT_TRASH:
             raise ValueError(
                 f"Invalid Deck Type: {deck_type.name} - Does not support modules as fixtures."

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

@@ -268,6 +268,13 @@ class FilesAddedUpdate:
     file_ids: list[str]
 
 
+@dataclasses.dataclass
+class AddressableAreaUsedUpdate:
+    """An update that says an addressable area has been used."""
+
+    addressable_area_name: str
+
+
 @dataclasses.dataclass
 class StateUpdate:
     """Represents an update to perform on engine state."""
@@ -308,6 +315,8 @@ class StateUpdate:
 
     files_added: FilesAddedUpdate | NoChangeType = NO_CHANGE
 
+    addressable_area_used: AddressableAreaUsedUpdate | NoChangeType = NO_CHANGE
+
     def append(self, other: Self) -> Self:
         """Apply another `StateUpdate` "on top of" this one.
 
@@ -334,7 +343,8 @@ def reduce(cls: typing.Type[Self], *args: Self) -> Self:
         return accumulator
 
     # These convenience functions let the caller avoid the boilerplate of constructing a
-    # complicated dataclass tree.
+    # complicated dataclass tree, and allow chaining.
+
     @typing.overload
     def set_pipette_location(
         self: Self, *, pipette_id: str, new_deck_point: DeckPoint
@@ -567,3 +577,10 @@ def set_absorbance_reader_lid(self: Self, module_id: str, is_lid_on: bool) -> Se
             module_id=module_id, is_lid_on=is_lid_on
         )
         return self
+
+    def set_addressable_area_used(self: Self, addressable_area_name: str) -> Self:
+        """Mark that an addressable area has been used. See `AddressableAreaUsedUpdate`."""
+        self.addressable_area_used = AddressableAreaUsedUpdate(
+            addressable_area_name=addressable_area_name
+        )
+        return self