feat(api): add a reload-labware command (#14963)

Adds a new command ReloadLabware, which allows dispatchers to change all
the details of a loaded labware except for the location. This is
primarily intended to allow getting a new labware offset that was not
added to the engine by the time this labware was loaded (though it can
technically do more, for symmetry).

This doesn't really change a whole lot of behavior and is well-supported
with testing. It's a prerequisite for #14940

Closes RSQ-29

Author: sfoster1

api/src/opentrons/protocol_engine/clients/sync_client.py

@@ -127,6 +127,19 @@ def load_labware(
 
         return cast(commands.LoadLabwareResult, result)
 
+    def reload_labware(
+        self,
+        labware_id: str,
+    ) -> commands.ReloadLabwareResult:
+        """Execute a ReloadLabware command and return the result."""
+        request = commands.ReloadLabwareCreate(
+            params=commands.ReloadLabwareParams(
+                labwareId=labware_id,
+            )
+        )
+        result = self._transport.execute_command(request=request)
+        return cast(commands.ReloadLabwareResult, result)
+
     # TODO (spp, 2022-12-14): https://opentrons.atlassian.net/browse/RLAB-237
     def move_labware(
         self,

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

@@ -120,6 +120,14 @@
     LoadLabwareCommandType,
 )
 
+from .reload_labware import (
+    ReloadLabware,
+    ReloadLabwareParams,
+    ReloadLabwareCreate,
+    ReloadLabwareResult,
+    ReloadLabwareCommandType,
+)
+
 from .load_liquid import (
     LoadLiquid,
     LoadLiquidParams,
@@ -402,6 +410,12 @@
     "LoadLabwareParams",
     "LoadLabwareResult",
     "LoadLabwareCommandType",
+    # reload labware command models
+    "ReloadLabware",
+    "ReloadLabwareCreate",
+    "ReloadLabwareParams",
+    "ReloadLabwareResult",
+    "ReloadLabwareCommandType",
     # load module command models
     "LoadModule",
     "LoadModuleCreate",

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

@@ -100,6 +100,14 @@
     LoadLabwareCommandType,
 )
 
+from .reload_labware import (
+    ReloadLabware,
+    ReloadLabwareParams,
+    ReloadLabwareCreate,
+    ReloadLabwareResult,
+    ReloadLabwareCommandType,
+)
+
 from .load_liquid import (
     LoadLiquid,
     LoadLiquidParams,
@@ -304,6 +312,7 @@
         Home,
         RetractAxis,
         LoadLabware,
+        ReloadLabware,
         LoadLiquid,
         LoadModule,
         LoadPipette,
@@ -368,6 +377,7 @@
     HomeParams,
     RetractAxisParams,
     LoadLabwareParams,
+    ReloadLabwareParams,
     LoadLiquidParams,
     LoadModuleParams,
     LoadPipetteParams,
@@ -431,6 +441,7 @@
     HomeCommandType,
     RetractAxisCommandType,
     LoadLabwareCommandType,
+    ReloadLabwareCommandType,
     LoadLiquidCommandType,
     LoadModuleCommandType,
     LoadPipetteCommandType,
@@ -494,6 +505,7 @@
         HomeCreate,
         RetractAxisCreate,
         LoadLabwareCreate,
+        ReloadLabwareCreate,
         LoadLiquidCreate,
         LoadModuleCreate,
         LoadPipetteCreate,
@@ -558,6 +570,7 @@
     HomeResult,
     RetractAxisResult,
     LoadLabwareResult,
+    ReloadLabwareResult,
     LoadLiquidResult,
     LoadModuleResult,
     LoadPipetteResult,

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

@@ -0,0 +1,86 @@
+"""Reload labware command request, result, and implementation models."""
+from __future__ import annotations
+from pydantic import BaseModel, Field
+from typing import TYPE_CHECKING, Optional, Type
+from typing_extensions import Literal
+
+from .command import AbstractCommandImpl, BaseCommand, BaseCommandCreate
+
+if TYPE_CHECKING:
+    from ..state import StateView
+    from ..execution import EquipmentHandler
+
+
+ReloadLabwareCommandType = Literal["reloadLabware"]
+
+
+class ReloadLabwareParams(BaseModel):
+    """Payload required to load a labware into a slot."""
+
+    labwareId: str = Field(
+        ..., description="The already-loaded labware instance to update."
+    )
+
+
+class ReloadLabwareResult(BaseModel):
+    """Result data from the execution of a LoadLabware command."""
+
+    labwareId: str = Field(
+        ...,
+        description="An ID to reference this labware in subsequent commands. Same as the one in the parameters.",
+    )
+    offsetId: Optional[str] = Field(
+        # Default `None` instead of `...` so this field shows up as non-required in
+        # OpenAPI. The server is allowed to omit it or make it null.
+        None,
+        description=(
+            "An ID referencing the labware offset that will apply"
+            " to the reloaded labware."
+            " This offset will be in effect until the labware is moved"
+            " with a `moveLabware` command."
+            " Null or undefined means no offset applies,"
+            " so the default of (0, 0, 0) will be used."
+        ),
+    )
+
+
+class ReloadLabwareImplementation(
+    AbstractCommandImpl[ReloadLabwareParams, ReloadLabwareResult]
+):
+    """Reload labware command implementation."""
+
+    def __init__(
+        self, equipment: EquipmentHandler, state_view: StateView, **kwargs: object
+    ) -> None:
+        self._equipment = equipment
+        self._state_view = state_view
+
+    async def execute(self, params: ReloadLabwareParams) -> ReloadLabwareResult:
+        """Reload the definition and calibration data for a specific labware."""
+        reloaded_labware = await self._equipment.reload_labware(
+            labware_id=params.labwareId,
+        )
+
+        return ReloadLabwareResult(
+            labwareId=params.labwareId,
+            offsetId=reloaded_labware.offsetId,
+        )
+
+
+class ReloadLabware(BaseCommand[ReloadLabwareParams, ReloadLabwareResult]):
+    """Reload labware command resource model."""
+
+    commandType: ReloadLabwareCommandType = "reloadLabware"
+    params: ReloadLabwareParams
+    result: Optional[ReloadLabwareResult]
+
+    _ImplementationCls: Type[ReloadLabwareImplementation] = ReloadLabwareImplementation
+
+
+class ReloadLabwareCreate(BaseCommandCreate[ReloadLabwareParams]):
+    """Reload labware command creation request."""
+
+    commandType: ReloadLabwareCommandType = "reloadLabware"
+    params: ReloadLabwareParams
+
+    _CommandCls: Type[ReloadLabware] = ReloadLabware

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

@@ -8,6 +8,7 @@
     LoadedPipetteData,
     LoadedModuleData,
     LoadedConfigureForVolumeData,
+    ReloadedLabwareData,
 )
 from .movement import MovementHandler
 from .gantry_mover import GantryMover
@@ -29,6 +30,7 @@
     "create_queue_worker",
     "EquipmentHandler",
     "LoadedLabwareData",
+    "ReloadedLabwareData",
     "LoadedPipetteData",
     "LoadedModuleData",
     "LoadedConfigureForVolumeData",

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

@@ -56,6 +56,14 @@ class LoadedLabwareData:
     offsetId: Optional[str]
 
 
+@dataclass(frozen=True)
+class ReloadedLabwareData:
+    """The result of a reload labware procedure."""
+
+    location: LabwareLocation
+    offsetId: Optional[str]
+
+
 @dataclass(frozen=True)
 class LoadedPipetteData:
     """The result of a load pipette procedure."""
@@ -171,6 +179,25 @@ async def load_labware(
             labware_id=labware_id, definition=definition, offsetId=offset_id
         )
 
+    async def reload_labware(self, labware_id: str) -> ReloadedLabwareData:
+        """Reload an already-loaded labware. This cannot change the labware location.
+
+        Args:
+            labware_id: The ID of the already-loaded labware.
+
+        Raises:
+            LabwareNotLoadedError: If `labware_id` does not reference a loaded labware.
+
+        """
+        location = self._state_store.labware.get_location(labware_id)
+        definition_uri = self._state_store.labware.get_definition_uri(labware_id)
+        offset_id = self.find_applicable_labware_offset_id(
+            labware_definition_uri=definition_uri,
+            labware_location=location,
+        )
+
+        return ReloadedLabwareData(location=location, offsetId=offset_id)
+
     async def load_pipette(
         self,
         pipette_name: PipetteNameType,

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

@@ -31,6 +31,7 @@
     Command,
     LoadLabwareResult,
     MoveLabwareResult,
+    ReloadLabwareResult,
 )
 from ..types import (
     DeckSlotLocation,
@@ -187,18 +188,27 @@ def _handle_command(self, command: Command) -> None:
             )
 
             self._state.definitions_by_uri[definition_uri] = command.result.definition
+            if isinstance(command.result, LoadLabwareResult):
+                location = command.params.location
+            else:
+                location = self._state.labware_by_id[command.result.labwareId].location
 
             self._state.labware_by_id[
                 command.result.labwareId
             ] = LoadedLabware.construct(
                 id=command.result.labwareId,
-                location=command.params.location,
+                location=location,
                 loadName=command.result.definition.parameters.loadName,
                 definitionUri=definition_uri,
                 offsetId=command.result.offsetId,
                 displayName=command.params.displayName,
             )
 
+        elif isinstance(command.result, ReloadLabwareResult):
+            labware_id = command.params.labwareId
+            new_offset_id = command.result.offsetId
+            self._state.labware_by_id[labware_id].offsetId = new_offset_id
+
         elif isinstance(command.result, MoveLabwareResult):
             labware_id = command.params.labwareId
             new_location = command.params.newLocation

api/tests/opentrons/protocol_engine/clients/test_sync_client.py

@@ -161,6 +161,30 @@ def test_load_labware(
     assert result == expected_result
 
 
+def test_reload_labware(
+    decoy: Decoy,
+    transport: ChildThreadTransport,
+    subject: SyncClient,
+) -> None:
+    """It should execute a reload labware command."""
+    expected_request = commands.ReloadLabwareCreate(
+        params=commands.ReloadLabwareParams(
+            labwareId="some-labware-id",
+        )
+    )
+
+    expected_result = commands.ReloadLabwareResult(
+        labwareId="some-labware-id", offsetId=None
+    )
+    decoy.when(transport.execute_command(request=expected_request)).then_return(
+        expected_result
+    )
+    result = subject.reload_labware(
+        labware_id="some-labware-id",
+    )
+    assert result == expected_result
+
+
 def test_load_module(
     decoy: Decoy,
     transport: ChildThreadTransport,