perf(robot-server): Discriminate command unions (#14286)

Closes RSS-128.

* Update command unions.
* Update stateless command unions.
* Work around Pydantic bug.
* Fix names and docstrings for union workarounds.
* Update tests.
* Update command schema.

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

@@ -1,6 +1,9 @@
 """Union types of concrete command definitions."""
 
 from typing import Union
+from typing_extensions import Annotated
+
+from pydantic import Field
 
 from . import heater_shaker
 from . import magnetic_module
@@ -284,66 +287,69 @@
     GetTipPresenceCommandType,
 )
 
-Command = Union[
-    Aspirate,
-    AspirateInPlace,
-    Comment,
-    Custom,
-    Dispense,
-    DispenseInPlace,
-    BlowOut,
-    BlowOutInPlace,
-    ConfigureForVolume,
-    ConfigureNozzleLayout,
-    DropTip,
-    DropTipInPlace,
-    Home,
-    RetractAxis,
-    LoadLabware,
-    LoadLiquid,
-    LoadModule,
-    LoadPipette,
-    MoveLabware,
-    MoveRelative,
-    MoveToCoordinates,
-    MoveToWell,
-    MoveToAddressableArea,
-    MoveToAddressableAreaForDropTip,
-    PrepareToAspirate,
-    WaitForResume,
-    WaitForDuration,
-    PickUpTip,
-    SavePosition,
-    SetRailLights,
-    TouchTip,
-    SetStatusBar,
-    VerifyTipPresence,
-    GetTipPresence,
-    heater_shaker.WaitForTemperature,
-    heater_shaker.SetTargetTemperature,
-    heater_shaker.DeactivateHeater,
-    heater_shaker.SetAndWaitForShakeSpeed,
-    heater_shaker.DeactivateShaker,
-    heater_shaker.OpenLabwareLatch,
-    heater_shaker.CloseLabwareLatch,
-    magnetic_module.Disengage,
-    magnetic_module.Engage,
-    temperature_module.SetTargetTemperature,
-    temperature_module.WaitForTemperature,
-    temperature_module.DeactivateTemperature,
-    thermocycler.SetTargetBlockTemperature,
-    thermocycler.WaitForBlockTemperature,
-    thermocycler.SetTargetLidTemperature,
-    thermocycler.WaitForLidTemperature,
-    thermocycler.DeactivateBlock,
-    thermocycler.DeactivateLid,
-    thermocycler.OpenLid,
-    thermocycler.CloseLid,
-    thermocycler.RunProfile,
-    calibration.CalibrateGripper,
-    calibration.CalibratePipette,
-    calibration.CalibrateModule,
-    calibration.MoveToMaintenancePosition,
+Command = Annotated[
+    Union[
+        Aspirate,
+        AspirateInPlace,
+        Comment,
+        Custom,
+        Dispense,
+        DispenseInPlace,
+        BlowOut,
+        BlowOutInPlace,
+        ConfigureForVolume,
+        ConfigureNozzleLayout,
+        DropTip,
+        DropTipInPlace,
+        Home,
+        RetractAxis,
+        LoadLabware,
+        LoadLiquid,
+        LoadModule,
+        LoadPipette,
+        MoveLabware,
+        MoveRelative,
+        MoveToCoordinates,
+        MoveToWell,
+        MoveToAddressableArea,
+        MoveToAddressableAreaForDropTip,
+        PrepareToAspirate,
+        WaitForResume,
+        WaitForDuration,
+        PickUpTip,
+        SavePosition,
+        SetRailLights,
+        TouchTip,
+        SetStatusBar,
+        VerifyTipPresence,
+        GetTipPresence,
+        heater_shaker.WaitForTemperature,
+        heater_shaker.SetTargetTemperature,
+        heater_shaker.DeactivateHeater,
+        heater_shaker.SetAndWaitForShakeSpeed,
+        heater_shaker.DeactivateShaker,
+        heater_shaker.OpenLabwareLatch,
+        heater_shaker.CloseLabwareLatch,
+        magnetic_module.Disengage,
+        magnetic_module.Engage,
+        temperature_module.SetTargetTemperature,
+        temperature_module.WaitForTemperature,
+        temperature_module.DeactivateTemperature,
+        thermocycler.SetTargetBlockTemperature,
+        thermocycler.WaitForBlockTemperature,
+        thermocycler.SetTargetLidTemperature,
+        thermocycler.WaitForLidTemperature,
+        thermocycler.DeactivateBlock,
+        thermocycler.DeactivateLid,
+        thermocycler.OpenLid,
+        thermocycler.CloseLid,
+        thermocycler.RunProfile,
+        calibration.CalibrateGripper,
+        calibration.CalibratePipette,
+        calibration.CalibrateModule,
+        calibration.MoveToMaintenancePosition,
+    ],
+    Field(discriminator="commandType"),
 ]
 
 CommandParams = Union[
@@ -471,66 +477,69 @@
     calibration.MoveToMaintenancePositionCommandType,
 ]
 
-CommandCreate = Union[
-    AspirateCreate,
-    AspirateInPlaceCreate,
-    CommentCreate,
-    ConfigureForVolumeCreate,
-    ConfigureNozzleLayoutCreate,
-    CustomCreate,
-    DispenseCreate,
-    DispenseInPlaceCreate,
-    BlowOutCreate,
-    BlowOutInPlaceCreate,
-    DropTipCreate,
-    DropTipInPlaceCreate,
-    HomeCreate,
-    RetractAxisCreate,
-    LoadLabwareCreate,
-    LoadLiquidCreate,
-    LoadModuleCreate,
-    LoadPipetteCreate,
-    MoveLabwareCreate,
-    MoveRelativeCreate,
-    MoveToCoordinatesCreate,
-    MoveToWellCreate,
-    MoveToAddressableAreaCreate,
-    MoveToAddressableAreaForDropTipCreate,
-    PrepareToAspirateCreate,
-    WaitForResumeCreate,
-    WaitForDurationCreate,
-    PickUpTipCreate,
-    SavePositionCreate,
-    SetRailLightsCreate,
-    TouchTipCreate,
-    SetStatusBarCreate,
-    VerifyTipPresenceCreate,
-    GetTipPresenceCreate,
-    heater_shaker.WaitForTemperatureCreate,
-    heater_shaker.SetTargetTemperatureCreate,
-    heater_shaker.DeactivateHeaterCreate,
-    heater_shaker.SetAndWaitForShakeSpeedCreate,
-    heater_shaker.DeactivateShakerCreate,
-    heater_shaker.OpenLabwareLatchCreate,
-    heater_shaker.CloseLabwareLatchCreate,
-    magnetic_module.DisengageCreate,
-    magnetic_module.EngageCreate,
-    temperature_module.SetTargetTemperatureCreate,
-    temperature_module.WaitForTemperatureCreate,
-    temperature_module.DeactivateTemperatureCreate,
-    thermocycler.SetTargetBlockTemperatureCreate,
-    thermocycler.WaitForBlockTemperatureCreate,
-    thermocycler.SetTargetLidTemperatureCreate,
-    thermocycler.WaitForLidTemperatureCreate,
-    thermocycler.DeactivateBlockCreate,
-    thermocycler.DeactivateLidCreate,
-    thermocycler.OpenLidCreate,
-    thermocycler.CloseLidCreate,
-    thermocycler.RunProfileCreate,
-    calibration.CalibrateGripperCreate,
-    calibration.CalibratePipetteCreate,
-    calibration.CalibrateModuleCreate,
-    calibration.MoveToMaintenancePositionCreate,
+CommandCreate = Annotated[
+    Union[
+        AspirateCreate,
+        AspirateInPlaceCreate,
+        CommentCreate,
+        ConfigureForVolumeCreate,
+        ConfigureNozzleLayoutCreate,
+        CustomCreate,
+        DispenseCreate,
+        DispenseInPlaceCreate,
+        BlowOutCreate,
+        BlowOutInPlaceCreate,
+        DropTipCreate,
+        DropTipInPlaceCreate,
+        HomeCreate,
+        RetractAxisCreate,
+        LoadLabwareCreate,
+        LoadLiquidCreate,
+        LoadModuleCreate,
+        LoadPipetteCreate,
+        MoveLabwareCreate,
+        MoveRelativeCreate,
+        MoveToCoordinatesCreate,
+        MoveToWellCreate,
+        MoveToAddressableAreaCreate,
+        MoveToAddressableAreaForDropTipCreate,
+        PrepareToAspirateCreate,
+        WaitForResumeCreate,
+        WaitForDurationCreate,
+        PickUpTipCreate,
+        SavePositionCreate,
+        SetRailLightsCreate,
+        TouchTipCreate,
+        SetStatusBarCreate,
+        VerifyTipPresenceCreate,
+        GetTipPresenceCreate,
+        heater_shaker.WaitForTemperatureCreate,
+        heater_shaker.SetTargetTemperatureCreate,
+        heater_shaker.DeactivateHeaterCreate,
+        heater_shaker.SetAndWaitForShakeSpeedCreate,
+        heater_shaker.DeactivateShakerCreate,
+        heater_shaker.OpenLabwareLatchCreate,
+        heater_shaker.CloseLabwareLatchCreate,
+        magnetic_module.DisengageCreate,
+        magnetic_module.EngageCreate,
+        temperature_module.SetTargetTemperatureCreate,
+        temperature_module.WaitForTemperatureCreate,
+        temperature_module.DeactivateTemperatureCreate,
+        thermocycler.SetTargetBlockTemperatureCreate,
+        thermocycler.WaitForBlockTemperatureCreate,
+        thermocycler.SetTargetLidTemperatureCreate,
+        thermocycler.WaitForLidTemperatureCreate,
+        thermocycler.DeactivateBlockCreate,
+        thermocycler.DeactivateLidCreate,
+        thermocycler.OpenLidCreate,
+        thermocycler.CloseLidCreate,
+        thermocycler.RunProfileCreate,
+        calibration.CalibrateGripperCreate,
+        calibration.CalibratePipetteCreate,
+        calibration.CalibrateModuleCreate,
+        calibration.MoveToMaintenancePositionCreate,
+    ],
+    Field(discriminator="commandType"),
 ]
 
 CommandResult = Union[

robot-server/robot_server/commands/router.py

@@ -27,6 +27,17 @@
 commands_router = APIRouter()
 
 
+class RequestModelWithStatelessCommandCreate(RequestModel[StatelessCommandCreate]):
+    """Equivalent to RequestModel[StatelessCommandCreate].
+
+    This works around a Pydantic v<2 bug where RequestModel[StatelessCommandCreate]
+    doesn't parse using the StatelessCommandCreate union discriminator.
+    https://github.com/pydantic/pydantic/issues/3782
+    """
+
+    data: StatelessCommandCreate
+
+
 class CommandNotFound(ErrorDetails):
     """An error returned if the given command cannot be found."""
 
@@ -50,7 +61,7 @@ class CommandNotFound(ErrorDetails):
     },
 )
 async def create_command(
-    request_body: RequestModel[StatelessCommandCreate],
+    request_body: RequestModelWithStatelessCommandCreate,
     waitUntilComplete: bool = Query(
         False,
         description=(

robot-server/robot_server/commands/stateless_commands.py

@@ -1,47 +1,57 @@
 """Command requests and responses allowed to be used with /commands."""
 from typing import Union
+from typing_extensions import Annotated
+
+from pydantic import Field
+
 from opentrons.protocol_engine import commands
 
-StatelessCommandCreate = Union[
-    commands.HomeCreate,
-    commands.SetRailLightsCreate,
-    commands.SetStatusBarCreate,
-    commands.magnetic_module.EngageCreate,
-    commands.magnetic_module.DisengageCreate,
-    commands.temperature_module.SetTargetTemperatureCreate,
-    commands.temperature_module.DeactivateTemperatureCreate,
-    commands.thermocycler.SetTargetBlockTemperatureCreate,
-    commands.thermocycler.SetTargetLidTemperatureCreate,
-    commands.thermocycler.DeactivateBlockCreate,
-    commands.thermocycler.DeactivateLidCreate,
-    commands.thermocycler.OpenLidCreate,
-    commands.thermocycler.CloseLidCreate,
-    commands.heater_shaker.SetTargetTemperatureCreate,
-    commands.heater_shaker.SetAndWaitForShakeSpeedCreate,
-    commands.heater_shaker.DeactivateHeaterCreate,
-    commands.heater_shaker.DeactivateShakerCreate,
-    commands.heater_shaker.OpenLabwareLatchCreate,
-    commands.heater_shaker.CloseLabwareLatchCreate,
+StatelessCommandCreate = Annotated[
+    Union[
+        commands.HomeCreate,
+        commands.SetRailLightsCreate,
+        commands.SetStatusBarCreate,
+        commands.magnetic_module.EngageCreate,
+        commands.magnetic_module.DisengageCreate,
+        commands.temperature_module.SetTargetTemperatureCreate,
+        commands.temperature_module.DeactivateTemperatureCreate,
+        commands.thermocycler.SetTargetBlockTemperatureCreate,
+        commands.thermocycler.SetTargetLidTemperatureCreate,
+        commands.thermocycler.DeactivateBlockCreate,
+        commands.thermocycler.DeactivateLidCreate,
+        commands.thermocycler.OpenLidCreate,
+        commands.thermocycler.CloseLidCreate,
+        commands.heater_shaker.SetTargetTemperatureCreate,
+        commands.heater_shaker.SetAndWaitForShakeSpeedCreate,
+        commands.heater_shaker.DeactivateHeaterCreate,
+        commands.heater_shaker.DeactivateShakerCreate,
+        commands.heater_shaker.OpenLabwareLatchCreate,
+        commands.heater_shaker.CloseLabwareLatchCreate,
+    ],
+    Field(discriminator="commandType"),
 ]
 
-StatelessCommand = Union[
-    commands.Home,
-    commands.SetRailLights,
-    commands.SetStatusBar,
-    commands.magnetic_module.Engage,
-    commands.magnetic_module.Disengage,
-    commands.temperature_module.SetTargetTemperature,
-    commands.temperature_module.DeactivateTemperature,
-    commands.thermocycler.SetTargetBlockTemperature,
-    commands.thermocycler.SetTargetLidTemperature,
-    commands.thermocycler.DeactivateBlock,
-    commands.thermocycler.DeactivateLid,
-    commands.thermocycler.OpenLid,
-    commands.thermocycler.CloseLid,
-    commands.heater_shaker.SetTargetTemperature,
-    commands.heater_shaker.SetAndWaitForShakeSpeed,
-    commands.heater_shaker.DeactivateHeater,
-    commands.heater_shaker.DeactivateShaker,
-    commands.heater_shaker.OpenLabwareLatch,
-    commands.heater_shaker.CloseLabwareLatch,
+StatelessCommand = Annotated[
+    Union[
+        commands.Home,
+        commands.SetRailLights,
+        commands.SetStatusBar,
+        commands.magnetic_module.Engage,
+        commands.magnetic_module.Disengage,
+        commands.temperature_module.SetTargetTemperature,
+        commands.temperature_module.DeactivateTemperature,
+        commands.thermocycler.SetTargetBlockTemperature,
+        commands.thermocycler.SetTargetLidTemperature,
+        commands.thermocycler.DeactivateBlock,
+        commands.thermocycler.DeactivateLid,
+        commands.thermocycler.OpenLid,
+        commands.thermocycler.CloseLid,
+        commands.heater_shaker.SetTargetTemperature,
+        commands.heater_shaker.SetAndWaitForShakeSpeed,
+        commands.heater_shaker.DeactivateHeater,
+        commands.heater_shaker.DeactivateShaker,
+        commands.heater_shaker.OpenLabwareLatch,
+        commands.heater_shaker.CloseLabwareLatch,
+    ],
+    Field(discriminator="commandType"),
 ]