feat(api): add stacker module context and protocol command skeletons (#17206)

This PR adds flex stacker support skeleton in the protocol engine. With
the changes in this PR, you can load a flex stacker into your protocol.

Author: ahiuchingau

api/src/opentrons/hardware_control/modules/__init__.py

@@ -62,7 +62,6 @@
     "AbsorbanceReader",
     "AbsorbanceReaderStatus",
     "AbsorbanceReaderDisconnectedError",
-    "ModuleDisconnectedCallback",
     "FlexStacker",
     "FlexStackerStatus",
     "PlatformState",

api/src/opentrons/protocol_api/core/common.py

@@ -11,6 +11,7 @@
     AbstractHeaterShakerCore,
     AbstractMagneticBlockCore,
     AbstractAbsorbanceReaderCore,
+    AbstractFlexStackerCore,
 )
 from .protocol import AbstractProtocol
 from .well import AbstractWellCore
@@ -27,5 +28,6 @@
 HeaterShakerCore = AbstractHeaterShakerCore
 MagneticBlockCore = AbstractMagneticBlockCore
 AbsorbanceReaderCore = AbstractAbsorbanceReaderCore
+FlexStackerCore = AbstractFlexStackerCore
 RobotCore = AbstractRobot
 ProtocolCore = AbstractProtocol[InstrumentCore, LabwareCore, ModuleCore]

api/src/opentrons/protocol_api/core/engine/module_core.py

@@ -37,6 +37,7 @@
     AbstractHeaterShakerCore,
     AbstractMagneticBlockCore,
     AbstractAbsorbanceReaderCore,
+    AbstractFlexStackerCore,
 )
 from .exceptions import InvalidMagnetEngageHeightError
 
@@ -692,3 +693,25 @@ def is_lid_on(self) -> bool:
             self.module_id
         )
         return abs_state.is_lid_on
+
+
+class FlexStackerCore(ModuleCore, AbstractFlexStackerCore):
+    """Flex Stacker core logic implementation for Python protocols."""
+
+    _sync_module_hardware: SynchronousAdapter[hw_modules.FlexStacker]
+
+    def retrieve(self) -> None:
+        """Retrieve a labware from the bottom of the Flex Stacker's stack."""
+        self._engine_client.execute_command(
+            cmd.flex_stacker.RetrieveParams(
+                moduleId=self.module_id,
+            )
+        )
+
+    def store(self) -> None:
+        """Store a labware at the bottom of the Flex Stacker's stack."""
+        self._engine_client.execute_command(
+            cmd.flex_stacker.StoreParams(
+                moduleId=self.module_id,
+            )
+        )

api/src/opentrons/protocol_api/core/module.py

@@ -379,3 +379,22 @@ def open_lid(self) -> None:
     @abstractmethod
     def is_lid_on(self) -> bool:
         """Return True if the Absorbance Reader's lid is currently closed."""
+
+
+class AbstractFlexStackerCore(AbstractModuleCore):
+    """Core control interface for an attached Flex Stacker."""
+
+    MODULE_TYPE: ClassVar = ModuleType.FLEX_STACKER
+
+    @abstractmethod
+    def get_serial_number(self) -> str:
+        """Get the module's unique hardware serial number."""
+
+    @abstractmethod
+    def retrieve(self) -> None:
+        """Release and return a labware at the bottom of the labware stack."""
+
+    @abstractmethod
+    def store(self) -> None:
+        """Store a labware at the bottom of the labware stack."""
+        pass

api/src/opentrons/protocol_api/module_contexts.py

@@ -29,6 +29,7 @@
     HeaterShakerCore,
     MagneticBlockCore,
     AbsorbanceReaderCore,
+    FlexStackerCore,
 )
 from .core.core_map import LoadedCoreMap
 from .core.engine import ENGINE_CORE_API_VERSION
@@ -1098,3 +1099,34 @@ def read(
         :returns: A dictionary of wavelengths to dictionary of values ordered by well name.
         """
         return self._core.read(filename=export_filename)
+
+
+class FlexStackerContext(ModuleContext):
+    """An object representing a connected Flex Stacker module.
+
+    It should not be instantiated directly; instead, it should be
+    created through :py:meth:`.ProtocolContext.load_module`.
+
+    .. versionadded:: 2.23
+    """
+
+    _core: FlexStackerCore
+
+    @property
+    @requires_version(2, 23)
+    def serial_number(self) -> str:
+        """Get the module's unique hardware serial number."""
+        return self._core.get_serial_number()
+
+    @requires_version(2, 23)
+    def retrieve(self) -> None:
+        """Release and return a labware at the bottom of the labware stack."""
+        self._core.retrieve()
+
+    @requires_version(2, 23)
+    def store(self, labware: Labware) -> None:
+        """Store a labware at the bottom of the labware stack.
+
+        :param labware: The labware object to store.
+        """
+        self._core.store()

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

@@ -14,6 +14,7 @@
 """
 
 from . import absorbance_reader
+from . import flex_stacker
 from . import heater_shaker
 from . import magnetic_module
 from . import temperature_module
@@ -614,6 +615,7 @@
     # hardware control command models
     # hardware module command bundles
     "absorbance_reader",
+    "flex_stacker",
     "heater_shaker",
     "magnetic_module",
     "temperature_module",

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

@@ -16,6 +16,7 @@
 from .movement_common import StallOrCollisionError
 
 from . import absorbance_reader
+from . import flex_stacker
 from . import heater_shaker
 from . import magnetic_module
 from . import temperature_module
@@ -430,6 +431,8 @@
         absorbance_reader.OpenLid,
         absorbance_reader.Initialize,
         absorbance_reader.ReadAbsorbance,
+        flex_stacker.Retrieve,
+        flex_stacker.Store,
         calibration.CalibrateGripper,
         calibration.CalibratePipette,
         calibration.CalibrateModule,
@@ -518,6 +521,8 @@
     absorbance_reader.OpenLidParams,
     absorbance_reader.InitializeParams,
     absorbance_reader.ReadAbsorbanceParams,
+    flex_stacker.RetrieveParams,
+    flex_stacker.StoreParams,
     calibration.CalibrateGripperParams,
     calibration.CalibratePipetteParams,
     calibration.CalibrateModuleParams,
@@ -604,6 +609,8 @@
     absorbance_reader.OpenLidCommandType,
     absorbance_reader.InitializeCommandType,
     absorbance_reader.ReadAbsorbanceCommandType,
+    flex_stacker.RetrieveCommandType,
+    flex_stacker.StoreCommandType,
     calibration.CalibrateGripperCommandType,
     calibration.CalibratePipetteCommandType,
     calibration.CalibrateModuleCommandType,
@@ -691,6 +698,8 @@
         absorbance_reader.OpenLidCreate,
         absorbance_reader.InitializeCreate,
         absorbance_reader.ReadAbsorbanceCreate,
+        flex_stacker.RetrieveCreate,
+        flex_stacker.StoreCreate,
         calibration.CalibrateGripperCreate,
         calibration.CalibratePipetteCreate,
         calibration.CalibrateModuleCreate,
@@ -786,6 +795,8 @@
     absorbance_reader.OpenLidResult,
     absorbance_reader.InitializeResult,
     absorbance_reader.ReadAbsorbanceResult,
+    flex_stacker.RetrieveResult,
+    flex_stacker.StoreResult,
     calibration.CalibrateGripperResult,
     calibration.CalibratePipetteResult,
     calibration.CalibrateModuleResult,

api/src/opentrons/protocol_engine/commands/flex_stacker/__init__.py

@@ -0,0 +1,33 @@
+"""Command models for Flex Stacker commands."""
+
+from .store import (
+    StoreCommandType,
+    StoreParams,
+    StoreResult,
+    Store,
+    StoreCreate,
+)
+
+from .retrieve import (
+    RetrieveCommandType,
+    RetrieveParams,
+    RetrieveResult,
+    Retrieve,
+    RetrieveCreate,
+)
+
+
+__all__ = [
+    # flexStacker/store
+    "StoreCommandType",
+    "StoreParams",
+    "StoreResult",
+    "Store",
+    "StoreCreate",
+    # flexStacker/retrieve
+    "RetrieveCommandType",
+    "RetrieveParams",
+    "RetrieveResult",
+    "Retrieve",
+    "RetrieveCreate",
+]

api/src/opentrons/protocol_engine/commands/flex_stacker/retrieve.py

@@ -0,0 +1,77 @@
+"""Command models to retrieve a labware from a Flex Stacker."""
+from __future__ import annotations
+from typing import Optional, Literal, TYPE_CHECKING
+from typing_extensions import Type
+
+from pydantic import BaseModel, Field
+
+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
+    from opentrons.protocol_engine.execution import EquipmentHandler
+
+RetrieveCommandType = Literal["flexStacker/retrieve"]
+
+
+class RetrieveParams(BaseModel):
+    """Input parameters for a labware retrieval command."""
+
+    moduleId: str = Field(
+        ...,
+        description="Unique ID of the Flex Stacker.",
+    )
+
+
+class RetrieveResult(BaseModel):
+    """Result data from a labware retrieval command."""
+
+
+class RetrieveImpl(AbstractCommandImpl[RetrieveParams, SuccessData[RetrieveResult]]):
+    """Implementation of a labware retrieval command."""
+
+    def __init__(
+        self,
+        state_view: StateView,
+        equipment: EquipmentHandler,
+        **kwargs: object,
+    ) -> None:
+        self._state_view = state_view
+        self._equipment = equipment
+
+    async def execute(self, params: RetrieveParams) -> SuccessData[RetrieveResult]:
+        """Execute the labware retrieval command."""
+        state_update = update_types.StateUpdate()
+        stacker_substate = self._state_view.modules.get_flex_stacker_substate(
+            module_id=params.moduleId
+        )
+
+        # Allow propagation of ModuleNotAttachedError.
+        stacker = self._equipment.get_module_hardware_api(stacker_substate.module_id)
+
+        if stacker is not None:
+            # TODO: get labware height from labware state view
+            await stacker.dispense_labware(labware_height=50.0)
+
+        return SuccessData(public=RetrieveResult(), state_update=state_update)
+
+
+class Retrieve(BaseCommand[RetrieveParams, RetrieveResult, ErrorOccurrence]):
+    """A command to retrieve a labware from a Flex Stacker."""
+
+    commandType: RetrieveCommandType = "flexStacker/retrieve"
+    params: RetrieveParams
+    result: Optional[RetrieveResult]
+
+    _ImplementationCls: Type[RetrieveImpl] = RetrieveImpl
+
+
+class RetrieveCreate(BaseCommandCreate[RetrieveParams]):
+    """A request to execute a Flex Stacker retrieve command."""
+
+    commandType: RetrieveCommandType = "flexStacker/retrieve"
+    params: RetrieveParams
+
+    _CommandCls: Type[Retrieve] = Retrieve

api/src/opentrons/protocol_engine/commands/flex_stacker/store.py

@@ -0,0 +1,78 @@
+"""Command models to retrieve a labware from a Flex Stacker."""
+from __future__ import annotations
+from typing import Optional, Literal, TYPE_CHECKING
+from typing_extensions import Type
+
+from pydantic import BaseModel, Field
+
+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
+    from opentrons.protocol_engine.execution import EquipmentHandler
+
+
+StoreCommandType = Literal["flexStacker/store"]
+
+
+class StoreParams(BaseModel):
+    """Input parameters for a labware storage command."""
+
+    moduleId: str = Field(
+        ...,
+        description="Unique ID of the flex stacker.",
+    )
+
+
+class StoreResult(BaseModel):
+    """Result data from a labware storage command."""
+
+
+class StoreImpl(AbstractCommandImpl[StoreParams, SuccessData[StoreResult]]):
+    """Implementation of a labware storage command."""
+
+    def __init__(
+        self,
+        state_view: StateView,
+        equipment: EquipmentHandler,
+        **kwargs: object,
+    ) -> None:
+        self._state_view = state_view
+        self._equipment = equipment
+
+    async def execute(self, params: StoreParams) -> SuccessData[StoreResult]:
+        """Execute the labware storage command."""
+        state_update = update_types.StateUpdate()
+        stacker_substate = self._state_view.modules.get_flex_stacker_substate(
+            module_id=params.moduleId
+        )
+
+        # Allow propagation of ModuleNotAttachedError.
+        stacker = self._equipment.get_module_hardware_api(stacker_substate.module_id)
+
+        if stacker is not None:
+            # TODO: get labware height from labware state view
+            await stacker.store_labware(labware_height=50.0)
+
+        return SuccessData(public=StoreResult(), state_update=state_update)
+
+
+class Store(BaseCommand[StoreParams, StoreResult, ErrorOccurrence]):
+    """A command to store a labware in a Flex Stacker."""
+
+    commandType: StoreCommandType = "flexStacker/store"
+    params: StoreParams
+    result: Optional[StoreResult]
+
+    _ImplementationCls: Type[StoreImpl] = StoreImpl
+
+
+class StoreCreate(BaseCommandCreate[StoreParams]):
+    """A request to execute a Flex Stacker store command."""
+
+    commandType: StoreCommandType = "flexStacker/store"
+    params: StoreParams
+
+    _CommandCls: Type[Store] = Store