opentrons 8.1.0__py2.py3-none-any.whl → 8.2.0__py2.py3-none-any.whl

opentrons/protocol_engine/commands/pick_up_tip.py

@@ -1,19 +1,17 @@
 """Pick up tip command request, result, and implementation models."""
 from __future__ import annotations
-from dataclasses import dataclass
 from opentrons_shared_data.errors import ErrorCodes
 from pydantic import Field
 from typing import TYPE_CHECKING, Optional, Type, Union
 from typing_extensions import Literal
 
-from opentrons.protocol_engine.errors.exceptions import TipNotAttachedError
 
-from ..errors import ErrorOccurrence
+from ..errors import ErrorOccurrence, PickUpTipTipNotAttachedError
 from ..resources import ModelUtils
-from ..types import DeckPoint
+from ..state import update_types
+from ..types import PickUpTipWellLocation, DeckPoint
 from .pipetting_common import (
     PipetteIdMixin,
-    WellLocationMixin,
     DestinationPositionResult,
 )
 from .command import (
@@ -25,17 +23,22 @@ from .command import (
 )
 
 if TYPE_CHECKING:
-    from ..state import StateView
+    from ..state.state import StateView
     from ..execution import MovementHandler, TipHandler
 
 
 PickUpTipCommandType = Literal["pickUpTip"]
 
 
-class PickUpTipParams(PipetteIdMixin, WellLocationMixin):
+class PickUpTipParams(PipetteIdMixin):
     """Payload needed to move a pipette to a specific well."""
 
-    pass
+    labwareId: str = Field(..., description="Identifier of labware to use.")
+    wellName: str = Field(..., description="Name of well to use in labware.")
+    wellLocation: PickUpTipWellLocation = Field(
+        default_factory=PickUpTipWellLocation,
+        description="Relative well location at which to pick up the tip.",
+    )
 
 
 class PickUpTipResult(DestinationPositionResult):
@@ -72,24 +75,19 @@ class TipPhysicallyMissingError(ErrorOccurrence):
     of the pipette.
     """
 
+    # The thing above about marking the tips as used makes it so that
+    # when the protocol is resumed and the Python Protocol API calls
+    # `get_next_tip()`, we'll move on to other tips as expected.
+
     isDefined: bool = True
     errorType: Literal["tipPhysicallyMissing"] = "tipPhysicallyMissing"
     errorCode: str = ErrorCodes.TIP_PICKUP_FAILED.value.code
-    detail: str = "No tip detected."
-
-
-@dataclass(frozen=True)
-class TipPhysicallyMissingErrorInternalData:
-    """Internal-to-ProtocolEngine data about a TipPhysicallyMissingError."""
-
-    pipette_id: str
-    labware_id: str
-    well_name: str
+    detail: str = "No Tip Detected"
 
 
 _ExecuteReturn = Union[
-    SuccessData[PickUpTipResult, None],
-    DefinedErrorData[TipPhysicallyMissingError, TipPhysicallyMissingErrorInternalData],
+    SuccessData[PickUpTipResult],
+    DefinedErrorData[TipPhysicallyMissingError],
 ]
 
 
@@ -111,19 +109,30 @@ class PickUpTipImplementation(AbstractCommandImpl[PickUpTipParams, _ExecuteRetur
 
     async def execute(
         self, params: PickUpTipParams
-    ) -> Union[SuccessData[PickUpTipResult, None], _ExecuteReturn]:
+    ) -> Union[SuccessData[PickUpTipResult], _ExecuteReturn]:
         """Move to and pick up a tip using the requested pipette."""
         pipette_id = params.pipetteId
         labware_id = params.labwareId
         well_name = params.wellName
-        well_location = params.wellLocation
 
+        state_update = update_types.StateUpdate()
+
+        well_location = self._state_view.geometry.convert_pick_up_tip_well_location(
+            well_location=params.wellLocation
+        )
         position = await self._movement.move_to_well(
             pipette_id=pipette_id,
             labware_id=labware_id,
             well_name=well_name,
             well_location=well_location,
         )
+        deck_point = DeckPoint.construct(x=position.x, y=position.y, z=position.z)
+        state_update.set_pipette_location(
+            pipette_id=pipette_id,
+            new_labware_id=labware_id,
+            new_well_name=well_name,
+            new_deck_point=deck_point,
+        )
 
         try:
             tip_geometry = await self._tip_handler.pick_up_tip(
@@ -131,7 +140,15 @@ class PickUpTipImplementation(AbstractCommandImpl[PickUpTipParams, _ExecuteRetur
                 labware_id=labware_id,
                 well_name=well_name,
             )
-        except TipNotAttachedError as e:
+        except PickUpTipTipNotAttachedError as e:
+            state_update_if_false_positive = update_types.StateUpdate()
+            state_update_if_false_positive.update_pipette_tip_state(
+                pipette_id=pipette_id,
+                tip_geometry=e.tip_geometry,
+            )
+            state_update.mark_tips_as_used(
+                pipette_id=pipette_id, labware_id=labware_id, well_name=well_name
+            )
             return DefinedErrorData(
                 public=TipPhysicallyMissingError(
                     id=self._model_utils.generate_id(),
@@ -144,21 +161,25 @@ class PickUpTipImplementation(AbstractCommandImpl[PickUpTipParams, _ExecuteRetur
                         )
                     ],
                 ),
-                private=TipPhysicallyMissingErrorInternalData(
-                    pipette_id=pipette_id,
-                    labware_id=labware_id,
-                    well_name=well_name,
-                ),
+                state_update=state_update,
+                state_update_if_false_positive=state_update_if_false_positive,
             )
         else:
+            state_update.update_pipette_tip_state(
+                pipette_id=pipette_id,
+                tip_geometry=tip_geometry,
+            )
+            state_update.mark_tips_as_used(
+                pipette_id=pipette_id, labware_id=labware_id, well_name=well_name
+            )
             return SuccessData(
                 public=PickUpTipResult(
                     tipVolume=tip_geometry.volume,
                     tipLength=tip_geometry.length,
                     tipDiameter=tip_geometry.diameter,
-                    position=DeckPoint(x=position.x, y=position.y, z=position.z),
+                    position=deck_point,
                 ),
-                private=None,
+                state_update=state_update,
             )
 
 

opentrons/protocol_engine/commands/pipetting_common.py

@@ -1,12 +1,11 @@
 """Common pipetting command base models."""
-from dataclasses import dataclass
 from opentrons_shared_data.errors import ErrorCodes
 from pydantic import BaseModel, Field
 from typing import Literal, Optional, Tuple, TypedDict
 
 from opentrons.protocol_engine.errors.error_occurrence import ErrorOccurrence
 
-from ..types import WellLocation, DeckPoint
+from ..types import WellLocation, LiquidHandlingWellLocation, DeckPoint
 
 
 class PipetteIdMixin(BaseModel):
@@ -69,6 +68,23 @@ class WellLocationMixin(BaseModel):
     )
 
 
+class LiquidHandlingWellLocationMixin(BaseModel):
+    """Mixin for command requests that take a location that's somewhere in a well."""
+
+    labwareId: str = Field(
+        ...,
+        description="Identifier of labware to use.",
+    )
+    wellName: str = Field(
+        ...,
+        description="Name of well to use in labware.",
+    )
+    wellLocation: LiquidHandlingWellLocation = Field(
+        default_factory=LiquidHandlingWellLocation,
+        description="Relative well location at which to perform the operation",
+    )
+
+
 class MovementMixin(BaseModel):
     """Mixin for command requests that move a pipette."""
 
@@ -114,6 +130,14 @@ class BaseLiquidHandlingResult(BaseModel):
 class DestinationPositionResult(BaseModel):
     """Mixin for command results that move a pipette."""
 
+    # todo(mm, 2024-08-02): Consider deprecating or redefining this.
+    #
+    # This is here because opentrons.protocol_engine needed it for internal bookkeeping
+    # and, at the time, we didn't have a way to do that without adding this to the
+    # public command results. Its usefulness to callers outside
+    # opentrons.protocol_engine is questionable because they would need to know which
+    # critical point is in play, and I think that can change depending on obscure
+    # things like labware quirks.
     position: DeckPoint = Field(
         DeckPoint(x=0, y=0, z=0),
         description=(
@@ -124,7 +148,12 @@ class DestinationPositionResult(BaseModel):
 
 
 class ErrorLocationInfo(TypedDict):
-    """Holds a retry location for in-place error recovery."""
+    """Holds a retry location for in-place error recovery.
+
+    This is appropriate to pass to a `moveToCoordinates` command,
+    assuming the pipette has not been configured with a different nozzle layout
+    in the meantime.
+    """
 
     retryLocation: Tuple[float, float, float]
 
@@ -149,14 +178,6 @@ class OverpressureError(ErrorOccurrence):
     errorInfo: ErrorLocationInfo
 
 
-@dataclass(frozen=True)
-class OverpressureErrorInternalData:
-    """Internal-to-ProtocolEngine data about an OverpressureError."""
-
-    position: DeckPoint
-    """Same meaning as DestinationPositionResult.position."""
-
-
 class LiquidNotFoundError(ErrorOccurrence):
     """Returned when no liquid is detected during the liquid probe process/move.
 
@@ -171,9 +192,19 @@ class LiquidNotFoundError(ErrorOccurrence):
     detail: str = ErrorCodes.PIPETTE_LIQUID_NOT_FOUND.value.detail
 
 
-@dataclass(frozen=True)
-class LiquidNotFoundErrorInternalData:
-    """Internal-to-ProtocolEngine data about a LiquidNotFoundError."""
+class TipPhysicallyAttachedError(ErrorOccurrence):
+    """Returned when sensors determine that a tip remains on the pipette after a drop attempt.
 
-    position: DeckPoint
-    """Same meaning as DestinationPositionResult.position."""
+    The pipette will act as if the tip was not dropped. So, you won't be able to pick
+    up a new tip without dropping the current one, and movement commands will assume
+    there is a tip hanging off the bottom of the pipette.
+    """
+
+    isDefined: bool = True
+
+    errorType: Literal["tipPhysicallyAttached"] = "tipPhysicallyAttached"
+
+    errorCode: str = ErrorCodes.TIP_DROP_FAILED.value.code
+    detail: str = ErrorCodes.TIP_DROP_FAILED.value.detail
+
+    errorInfo: ErrorLocationInfo

opentrons/protocol_engine/commands/prepare_to_aspirate.py

@@ -1,18 +1,28 @@
 """Prepare to aspirate command request, result, and implementation models."""
 
 from __future__ import annotations
+from opentrons_shared_data.errors.exceptions import PipetteOverpressureError
 from pydantic import BaseModel
-from typing import TYPE_CHECKING, Optional, Type
+from typing import TYPE_CHECKING, Optional, Type, Union
 from typing_extensions import Literal
 
 from .pipetting_common import (
+    OverpressureError,
     PipetteIdMixin,
 )
-from .command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
+from .command import (
+    AbstractCommandImpl,
+    BaseCommand,
+    BaseCommandCreate,
+    DefinedErrorData,
+    SuccessData,
+)
 from ..errors.error_occurrence import ErrorOccurrence
 
 if TYPE_CHECKING:
-    from ..execution.pipetting import PipettingHandler
+    from ..execution import PipettingHandler, GantryMover
+    from ..resources import ModelUtils
+
 
 PrepareToAspirateCommandType = Literal["prepareToAspirate"]
 
@@ -29,25 +39,62 @@ class PrepareToAspirateResult(BaseModel):
     pass
 
 
+_ExecuteReturn = Union[
+    SuccessData[PrepareToAspirateResult],
+    DefinedErrorData[OverpressureError],
+]
+
+
 class PrepareToAspirateImplementation(
-    AbstractCommandImpl[
-        PrepareToAspirateParams, SuccessData[PrepareToAspirateResult, None]
-    ]
+    AbstractCommandImpl[PrepareToAspirateParams, _ExecuteReturn]
 ):
     """Prepare for aspirate command implementation."""
 
-    def __init__(self, pipetting: PipettingHandler, **kwargs: object) -> None:
+    def __init__(
+        self,
+        pipetting: PipettingHandler,
+        model_utils: ModelUtils,
+        gantry_mover: GantryMover,
+        **kwargs: object,
+    ) -> None:
         self._pipetting_handler = pipetting
+        self._model_utils = model_utils
+        self._gantry_mover = gantry_mover
 
-    async def execute(
-        self, params: PrepareToAspirateParams
-    ) -> SuccessData[PrepareToAspirateResult, None]:
+    async def execute(self, params: PrepareToAspirateParams) -> _ExecuteReturn:
         """Prepare the pipette to aspirate."""
-        await self._pipetting_handler.prepare_for_aspirate(
-            pipette_id=params.pipetteId,
-        )
-
-        return SuccessData(public=PrepareToAspirateResult(), private=None)
+        current_position = await self._gantry_mover.get_position(params.pipetteId)
+        try:
+            await self._pipetting_handler.prepare_for_aspirate(
+                pipette_id=params.pipetteId,
+            )
+        except PipetteOverpressureError as e:
+            return DefinedErrorData(
+                public=OverpressureError(
+                    id=self._model_utils.generate_id(),
+                    createdAt=self._model_utils.get_timestamp(),
+                    wrappedErrors=[
+                        ErrorOccurrence.from_failed(
+                            id=self._model_utils.generate_id(),
+                            createdAt=self._model_utils.get_timestamp(),
+                            error=e,
+                        )
+                    ],
+                    errorInfo=(
+                        {
+                            "retryLocation": (
+                                current_position.x,
+                                current_position.y,
+                                current_position.z,
+                            )
+                        }
+                    ),
+                ),
+            )
+        else:
+            return SuccessData(
+                public=PrepareToAspirateResult(),
+            )
 
 
 class PrepareToAspirate(

opentrons/protocol_engine/commands/reload_labware.py

@@ -6,9 +6,10 @@ from typing_extensions import Literal
 
 from .command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
 from ..errors.error_occurrence import ErrorOccurrence
+from ..state.update_types import StateUpdate
 
 if TYPE_CHECKING:
-    from ..state import StateView
+    from ..state.state import StateView
     from ..execution import EquipmentHandler
 
 
@@ -46,7 +47,7 @@ class ReloadLabwareResult(BaseModel):
 
 
 class ReloadLabwareImplementation(
-    AbstractCommandImpl[ReloadLabwareParams, SuccessData[ReloadLabwareResult, None]]
+    AbstractCommandImpl[ReloadLabwareParams, SuccessData[ReloadLabwareResult]]
 ):
     """Reload labware command implementation."""
 
@@ -58,18 +59,26 @@ class ReloadLabwareImplementation(
 
     async def execute(
         self, params: ReloadLabwareParams
-    ) -> SuccessData[ReloadLabwareResult, None]:
+    ) -> SuccessData[ReloadLabwareResult]:
         """Reload the definition and calibration data for a specific labware."""
         reloaded_labware = await self._equipment.reload_labware(
             labware_id=params.labwareId,
         )
 
+        state_update = StateUpdate()
+
+        state_update.set_labware_location(
+            labware_id=params.labwareId,
+            new_location=reloaded_labware.location,
+            new_offset_id=reloaded_labware.offsetId,
+        )
+
         return SuccessData(
             public=ReloadLabwareResult(
                 labwareId=params.labwareId,
                 offsetId=reloaded_labware.offsetId,
             ),
-            private=None,
+            state_update=state_update,
         )
 
 

opentrons/protocol_engine/commands/retract_axis.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel, Field
 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
@@ -38,7 +39,7 @@ class RetractAxisResult(BaseModel):
 
 
 class RetractAxisImplementation(
-    AbstractCommandImpl[RetractAxisParams, SuccessData[RetractAxisResult, None]]
+    AbstractCommandImpl[RetractAxisParams, SuccessData[RetractAxisResult]]
 ):
     """Retract Axis command implementation."""
 
@@ -47,10 +48,12 @@ class RetractAxisImplementation(
 
     async def execute(
         self, params: RetractAxisParams
-    ) -> SuccessData[RetractAxisResult, None]:
+    ) -> SuccessData[RetractAxisResult]:
         """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(), state_update=state_update)
 
 
 class RetractAxis(BaseCommand[RetractAxisParams, RetractAxisResult, ErrorOccurrence]):

opentrons/protocol_engine/commands/save_position.py

@@ -46,7 +46,7 @@ class SavePositionResult(BaseModel):
 
 
 class SavePositionImplementation(
-    AbstractCommandImpl[SavePositionParams, SuccessData[SavePositionResult, None]]
+    AbstractCommandImpl[SavePositionParams, SuccessData[SavePositionResult]]
 ):
     """Save position command implementation."""
 
@@ -61,7 +61,7 @@ class SavePositionImplementation(
 
     async def execute(
         self, params: SavePositionParams
-    ) -> SuccessData[SavePositionResult, None]:
+    ) -> SuccessData[SavePositionResult]:
         """Check the requested pipette's current position."""
         position_id = self._model_utils.ensure_id(params.positionId)
         fail_on_not_homed = (
@@ -76,7 +76,6 @@ class SavePositionImplementation(
                 positionId=position_id,
                 position=DeckPoint(x=x, y=y, z=z),
             ),
-            private=None,
         )
 
 

opentrons/protocol_engine/commands/set_rail_lights.py

@@ -29,7 +29,7 @@ class SetRailLightsResult(BaseModel):
 
 
 class SetRailLightsImplementation(
-    AbstractCommandImpl[SetRailLightsParams, SuccessData[SetRailLightsResult, None]]
+    AbstractCommandImpl[SetRailLightsParams, SuccessData[SetRailLightsResult]]
 ):
     """setRailLights command implementation."""
 
@@ -38,10 +38,12 @@ class SetRailLightsImplementation(
 
     async def execute(
         self, params: SetRailLightsParams
-    ) -> SuccessData[SetRailLightsResult, None]:
+    ) -> SuccessData[SetRailLightsResult]:
         """Dispatch a set lights command setting the state of the rail lights."""
         await self._rail_lights.set_rail_lights(params.on)
-        return SuccessData(public=SetRailLightsResult(), private=None)
+        return SuccessData(
+            public=SetRailLightsResult(),
+        )
 
 
 class SetRailLights(

opentrons/protocol_engine/commands/set_status_bar.py

@@ -49,7 +49,7 @@ class SetStatusBarResult(BaseModel):
 
 
 class SetStatusBarImplementation(
-    AbstractCommandImpl[SetStatusBarParams, SuccessData[SetStatusBarResult, None]]
+    AbstractCommandImpl[SetStatusBarParams, SuccessData[SetStatusBarResult]]
 ):
     """setStatusBar command implementation."""
 
@@ -58,12 +58,14 @@ class SetStatusBarImplementation(
 
     async def execute(
         self, params: SetStatusBarParams
-    ) -> SuccessData[SetStatusBarResult, None]:
+    ) -> SuccessData[SetStatusBarResult]:
         """Execute the setStatusBar command."""
         if not self._status_bar.status_bar_should_not_be_changed():
             state = _animation_to_status_bar_state(params.animation)
             await self._status_bar.set_status_bar(state)
-        return SuccessData(public=SetStatusBarResult(), private=None)
+        return SuccessData(
+            public=SetStatusBarResult(),
+        )
 
 
 class SetStatusBar(

opentrons/protocol_engine/commands/temperature_module/deactivate.py

@@ -9,7 +9,7 @@ from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, Succe
 from ...errors.error_occurrence import ErrorOccurrence
 
 if TYPE_CHECKING:
-    from opentrons.protocol_engine.state import StateView
+    from opentrons.protocol_engine.state.state import StateView
     from opentrons.protocol_engine.execution import EquipmentHandler
 
 DeactivateTemperatureCommandType = Literal["temperatureModule/deactivate"]
@@ -27,7 +27,7 @@ class DeactivateTemperatureResult(BaseModel):
 
 class DeactivateTemperatureImpl(
     AbstractCommandImpl[
-        DeactivateTemperatureParams, SuccessData[DeactivateTemperatureResult, None]
+        DeactivateTemperatureParams, SuccessData[DeactivateTemperatureResult]
     ]
 ):
     """Execution implementation of a Temperature Module's deactivate command."""
@@ -43,7 +43,7 @@ class DeactivateTemperatureImpl(
 
     async def execute(
         self, params: DeactivateTemperatureParams
-    ) -> SuccessData[DeactivateTemperatureResult, None]:
+    ) -> SuccessData[DeactivateTemperatureResult]:
         """Deactivate a Temperature Module."""
         # Allow propagation of ModuleNotLoadedError and WrongModuleTypeError.
         module_substate = self._state_view.modules.get_temperature_module_substate(
@@ -57,7 +57,9 @@ class DeactivateTemperatureImpl(
 
         if temp_hardware_module is not None:
             await temp_hardware_module.deactivate()
-        return SuccessData(public=DeactivateTemperatureResult(), private=None)
+        return SuccessData(
+            public=DeactivateTemperatureResult(),
+        )
 
 
 class DeactivateTemperature(

opentrons/protocol_engine/commands/temperature_module/set_target_temperature.py

@@ -9,7 +9,7 @@ from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, Succe
 from ...errors.error_occurrence import ErrorOccurrence
 
 if TYPE_CHECKING:
-    from opentrons.protocol_engine.state import StateView
+    from opentrons.protocol_engine.state.state import StateView
     from opentrons.protocol_engine.execution import EquipmentHandler
 
 SetTargetTemperatureCommandType = Literal["temperatureModule/setTargetTemperature"]
@@ -34,7 +34,7 @@ class SetTargetTemperatureResult(BaseModel):
 
 class SetTargetTemperatureImpl(
     AbstractCommandImpl[
-        SetTargetTemperatureParams, SuccessData[SetTargetTemperatureResult, None]
+        SetTargetTemperatureParams, SuccessData[SetTargetTemperatureResult]
     ]
 ):
     """Execution implementation of a Temperature Module's set temperature command."""
@@ -50,7 +50,7 @@ class SetTargetTemperatureImpl(
 
     async def execute(
         self, params: SetTargetTemperatureParams
-    ) -> SuccessData[SetTargetTemperatureResult, None]:
+    ) -> SuccessData[SetTargetTemperatureResult]:
         """Set a Temperature Module's target temperature."""
         # Allow propagation of ModuleNotLoadedError and WrongModuleTypeError.
         module_substate = self._state_view.modules.get_temperature_module_substate(
@@ -69,7 +69,6 @@ class SetTargetTemperatureImpl(
             await temp_hardware_module.start_set_temperature(celsius=validated_temp)
         return SuccessData(
             public=SetTargetTemperatureResult(targetTemperature=validated_temp),
-            private=None,
         )
 
 

opentrons/protocol_engine/commands/temperature_module/wait_for_temperature.py

@@ -9,7 +9,7 @@ from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, Succe
 from ...errors.error_occurrence import ErrorOccurrence
 
 if TYPE_CHECKING:
-    from opentrons.protocol_engine.state import StateView
+    from opentrons.protocol_engine.state.state import StateView
     from opentrons.protocol_engine.execution import EquipmentHandler
 
 WaitForTemperatureCommandType = Literal["temperatureModule/waitForTemperature"]
@@ -35,9 +35,7 @@ class WaitForTemperatureResult(BaseModel):
 
 
 class WaitForTemperatureImpl(
-    AbstractCommandImpl[
-        WaitForTemperatureParams, SuccessData[WaitForTemperatureResult, None]
-    ]
+    AbstractCommandImpl[WaitForTemperatureParams, SuccessData[WaitForTemperatureResult]]
 ):
     """Execution implementation of Temperature Module's wait for temperature command."""
 
@@ -52,7 +50,7 @@ class WaitForTemperatureImpl(
 
     async def execute(
         self, params: WaitForTemperatureParams
-    ) -> SuccessData[WaitForTemperatureResult, None]:
+    ) -> SuccessData[WaitForTemperatureResult]:
         """Wait for a Temperature Module's target temperature."""
         # Allow propagation of ModuleNotLoadedError and WrongModuleTypeError.
         module_substate = self._state_view.modules.get_temperature_module_substate(
@@ -74,7 +72,9 @@ class WaitForTemperatureImpl(
             await temp_hardware_module.await_temperature(
                 awaiting_temperature=target_temp
             )
-        return SuccessData(public=WaitForTemperatureResult(), private=None)
+        return SuccessData(
+            public=WaitForTemperatureResult(),
+        )
 
 
 class WaitForTemperature(

opentrons/protocol_engine/commands/thermocycler/__init__.py

@@ -73,6 +73,16 @@ from .run_profile import (
     RunProfileCreate,
 )
 
+from .run_extended_profile import (
+    RunExtendedProfileCommandType,
+    RunExtendedProfileParams,
+    RunExtendedProfileResult,
+    RunExtendedProfile,
+    RunExtendedProfileCreate,
+    ProfileCycle,
+    ProfileStep,
+)
+
 
 __all__ = [
     # Set target block temperature command models
@@ -130,4 +140,13 @@ __all__ = [
     "RunProfileResult",
     "RunProfile",
     "RunProfileCreate",
+    # Run extended profile command models.
+    "RunExtendedProfileCommandType",
+    "RunExtendedProfileParams",
+    "RunExtendedProfileStepParams",
+    "RunExtendedProfileResult",
+    "RunExtendedProfile",
+    "RunExtendedProfileCreate",
+    "ProfileCycle",
+    "ProfileStep",
 ]