Genericize command requests and responses

robot-server/robot_server/service/session/command_execution/command.py

@@ -1,9 +1,9 @@
 from datetime import datetime
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Optional, Generic, TypeVar
 
 from robot_server.service.session.models.command import (
-    CommandStatus, CommandResultType, RequestTypes)
+    CommandStatus, RequestTypes)
 from robot_server.service.session.models.common import (
     IdentifierType, create_identifier)
 from opentrons.util.helpers import utc_now
@@ -15,12 +15,15 @@ class CommandMeta:
     created_at: datetime = field(default_factory=utc_now)
 
 
+ResultTypeT = TypeVar("ResultTypeT")
+
+
 @dataclass(frozen=True)
-class CommandResult:
+class CommandResult(Generic[ResultTypeT]):
     started_at: datetime
     completed_at: datetime
     status: CommandStatus = CommandStatus.executed
-    data: Optional[CommandResultType] = None
+    data: Optional[ResultTypeT] = None
 
 
 @dataclass(frozen=True)

robot-server/robot_server/service/session/models/command.py

@@ -2,22 +2,22 @@
 from enum import Enum
 import typing
 
+from typing_extensions import Literal
+from pydantic import BaseModel, Field
+from pydantic.generics import GenericModel
+
 from opentrons_shared_data.labware.dev_types import LabwareDefinition
 from opentrons_shared_data.pipette.dev_types import PipetteName
-from typing_extensions import Literal
+from opentrons.util.helpers import utc_now
 
-from robot_server.service.session.models.command_definitions import \
-    CommandDefinitionType, RobotCommand, ProtocolCommand, EquipmentCommand, \
-    PipetteCommand, CalibrationCommand, DeckCalibrationCommand, \
-    CheckCalibrationCommand
+from robot_server.service.session.models.command_definitions import (
+    ProtocolCommand, EquipmentCommand, PipetteCommand, CalibrationCommand,
+    DeckCalibrationCommand, CheckCalibrationCommand, CommandDefinitionType)
 from robot_server.service.session.models.common import (
     EmptyModel, JogPosition, IdentifierType, OffsetVector)
-from pydantic import BaseModel, Field
-from pydantic.generics import GenericModel
 from robot_server.service.legacy.models.control import Mount
 from robot_server.service.json_api import (
     ResponseModel, RequestModel, ResponseDataModel)
-from opentrons.util.helpers import utc_now
 
 
 class LoadLabwareRequestData(BaseModel):
@@ -112,128 +112,213 @@ class CommandStatus(str, Enum):
     failed = "failed"
 
 
+CommandT = typing.TypeVar('CommandT', bound=CommandDefinitionType)
 RequestDataT = typing.TypeVar('RequestDataT', bound=BaseModel)
+ResponseDataT = typing.TypeVar('ResponseDataT')
 
 
-class SessionCommandRequest(GenericModel, typing.Generic[RequestDataT]):
+class SessionCommandRequest(
+    GenericModel,
+    typing.Generic[CommandT, RequestDataT, ResponseDataT]
+):
     """A session command"""
-    command: CommandDefinitionType = Field(
+    command: CommandT = Field(
         ...,
         description="The command description")
+    data: RequestDataT = Field(
+        ...,
+        description="The command data"
+    )
+
+    def make_response(
+            self,
+            identifier: str,
+            status: CommandStatus,
+            created_at: datetime,
+            started_at: typing.Optional[datetime],
+            completed_at: typing.Optional[datetime],
+            result: typing.Optional[ResponseDataT]
+    ) -> 'SessionCommandResponse[CommandT, RequestDataT, ResponseDataT]':
+        return SessionCommandResponse(
+            command=self.command,
+            data=self.data,
+            id=identifier,
+            status=status,
+            createdAt=created_at,
+            startedAt=started_at,
+            completedAt=completed_at,
+            result=result)
+
+
+class SessionCommandResponse(
+    ResponseDataModel,
+    GenericModel,
+    typing.Generic[CommandT, RequestDataT, ResponseDataT]
+):
+    """A session command response"""
+    command: CommandT
     data: RequestDataT
+    status: CommandStatus
+    createdAt: datetime = Field(..., default_factory=utc_now)
+    startedAt: typing.Optional[datetime]
+    completedAt: typing.Optional[datetime]
+    result: typing.Optional[ResponseDataT] = None
+
+
+# The command definitions requiring no data and result types.
+CommandsEmptyData = Literal[
+    ProtocolCommand.start_run,
+    ProtocolCommand.start_simulate,
+    ProtocolCommand.cancel,
+    ProtocolCommand.pause,
+    ProtocolCommand.resume,
+    CalibrationCommand.load_labware,
+    CalibrationCommand.move_to_tip_rack,
+    CalibrationCommand.move_to_point_one,
+    CalibrationCommand.move_to_deck,
+    CalibrationCommand.move_to_reference_point,
+    CalibrationCommand.pick_up_tip,
+    CalibrationCommand.confirm_tip_attached,
+    CalibrationCommand.invalidate_tip,
+    CalibrationCommand.save_offset,
+    CalibrationCommand.exit,
+    CalibrationCommand.invalidate_last_action,
+    DeckCalibrationCommand.move_to_point_two,
+    DeckCalibrationCommand.move_to_point_three,
+    CheckCalibrationCommand.compare_point,
+    CheckCalibrationCommand.switch_pipette,
+    CheckCalibrationCommand.return_tip,
+    CheckCalibrationCommand.transition
+]
+
+
+class SimpleCommandRequest(
+    SessionCommandRequest[CommandsEmptyData,
+                          EmptyModel,
+                          EmptyModel]):
+    """A command containing no data and result type"""
+    pass
+
+
+SimpleCommandResponse = SessionCommandResponse[
+    CommandsEmptyData,
+    EmptyModel,
+    EmptyModel
+]
 
 
-class RobotCommandRequest(SessionCommandRequest[EmptyModel]):
-    command: Literal[
-        RobotCommand.home_all_motors,
-        RobotCommand.home_pipette,
-        RobotCommand.toggle_lights
-    ]
+LoadLabwareRequest = SessionCommandRequest[
+    Literal[EquipmentCommand.load_labware],
+    LoadLabwareRequestData,
+    LoadLabwareResponseData
+]
 
 
-class ProtocolCommandRequest(SessionCommandRequest[EmptyModel]):
-    command: Literal[
-        ProtocolCommand.start_run,
-        ProtocolCommand.start_simulate,
-        ProtocolCommand.cancel,
-        ProtocolCommand.pause,
-        ProtocolCommand.resume
-    ]
+LoadLabwareResponse = SessionCommandResponse[
+    Literal[EquipmentCommand.load_labware],
+    LoadLabwareRequestData,
+    LoadLabwareResponseData
+]
 
 
-class LoadLabwareRequest(SessionCommandRequest[LoadLabwareRequestData]):
-    command: Literal[EquipmentCommand.load_labware]
+LoadInstrumentRequest = SessionCommandRequest[
+    Literal[EquipmentCommand.load_instrument],
+    LoadInstrumentRequestData,
+    LoadInstrumentResponseData
+]
 
 
-class LoadInstrumentRequest(SessionCommandRequest[LoadInstrumentRequestData]):
-    command: Literal[EquipmentCommand.load_instrument]
+LoadInstrumentResponse = SessionCommandResponse[
+    Literal[EquipmentCommand.load_instrument],
+    LoadInstrumentRequestData,
+    LoadInstrumentResponseData
+]
 
 
-class LiquidRequest(SessionCommandRequest[LiquidRequestData]):
-    command: Literal[
-        PipetteCommand.aspirate,
-        PipetteCommand.dispense
-    ]
+LiquidRequest = SessionCommandRequest[
+    Literal[PipetteCommand.aspirate,
+            PipetteCommand.dispense],
+    LiquidRequestData,
+    EmptyModel
+]
 
 
-class TipRequest(SessionCommandRequest[PipetteRequestDataBase]):
-    command: Literal[
-        PipetteCommand.drop_tip,
-        PipetteCommand.pick_up_tip
-    ]
+LiquidResponse = SessionCommandResponse[
+    Literal[PipetteCommand.aspirate,
+            PipetteCommand.dispense],
+    LiquidRequestData,
+    EmptyModel
+]
 
 
-class CalibrationRequest(SessionCommandRequest[EmptyModel]):
-    command: Literal[
-        CalibrationCommand.load_labware,
-        CalibrationCommand.move_to_tip_rack,
-        CalibrationCommand.move_to_point_one,
-        CalibrationCommand.move_to_deck,
-        CalibrationCommand.move_to_reference_point,
-        CalibrationCommand.pick_up_tip,
-        CalibrationCommand.confirm_tip_attached,
-        CalibrationCommand.invalidate_tip,
-        CalibrationCommand.save_offset,
-        CalibrationCommand.exit,
-        CalibrationCommand.invalidate_last_action,
-    ]
+TipRequest = SessionCommandRequest[
+    Literal[PipetteCommand.drop_tip,
+            PipetteCommand.pick_up_tip],
+    PipetteRequestDataBase,
+    EmptyModel
+]
 
 
-class JogRequest(SessionCommandRequest[JogPosition]):
-    command: Literal[CalibrationCommand.jog]
+TipResponse = SessionCommandResponse[
+    Literal[PipetteCommand.drop_tip,
+            PipetteCommand.pick_up_tip],
+    PipetteRequestDataBase,
+    EmptyModel
+]
 
 
-class SetHasCalibrationBlockRequestM(
-    SessionCommandRequest[SetHasCalibrationBlockRequestData]
-):
-    command: Literal[CalibrationCommand.set_has_calibration_block]
+JogRequest = SessionCommandRequest[
+    Literal[CalibrationCommand.jog],
+    JogPosition,
+    EmptyModel
+]
 
 
-class DeckCalibrationCommandRequest(SessionCommandRequest[EmptyModel]):
-    command: Literal[
-        DeckCalibrationCommand.move_to_point_two,
-        DeckCalibrationCommand.move_to_point_three
-    ]
+JogResponse = SessionCommandResponse[
+    Literal[CalibrationCommand.jog],
+    JogPosition,
+    EmptyModel
+]
 
 
-class CheckCalibrationCommandRequest(SessionCommandRequest[EmptyModel]):
-    command: Literal[
-        CheckCalibrationCommand.compare_point,
-        CheckCalibrationCommand.switch_pipette,
-        CheckCalibrationCommand.return_tip,
-        CheckCalibrationCommand.transition
-    ]
+SetHasCalibrationBlockRequest = SessionCommandRequest[
+    Literal[CalibrationCommand.set_has_calibration_block],
+    SetHasCalibrationBlockRequestData,
+    EmptyModel
+]
 
 
-class SessionCommand(ResponseDataModel):
-    """A session command response"""
-    command: CommandDefinitionType
-    data: CommandDataType
-    status: CommandStatus
-    createdAt: datetime = Field(..., default_factory=utc_now)
-    startedAt: typing.Optional[datetime]
-    completedAt: typing.Optional[datetime]
-    result: typing.Optional[CommandResultType] = None
+SetHasCalibrationBlockResponse = SessionCommandResponse[
+    Literal[
+        CalibrationCommand.set_has_calibration_block],
+    SetHasCalibrationBlockRequestData,
+    EmptyModel]
 
 
 RequestTypes = typing.Union[
-    RobotCommandRequest,
-    ProtocolCommandRequest,
+    SimpleCommandRequest,
     LoadLabwareRequest,
     LoadInstrumentRequest,
     LiquidRequest,
     TipRequest,
-    CalibrationRequest,
     JogRequest,
-    SetHasCalibrationBlockRequestM,
-    DeckCalibrationCommandRequest,
-    CheckCalibrationCommandRequest
+    SetHasCalibrationBlockRequest,
+]
+
+ResponseTypes = typing.Union[
+    SimpleCommandResponse,
+    LoadLabwareResponse,
+    LoadInstrumentResponse,
+    LiquidResponse,
+    TipResponse,
+    JogResponse,
+    SetHasCalibrationBlockResponse,
 ]
 
 # Session command requests/responses
 CommandRequest = RequestModel[
     RequestTypes
 ]
 CommandResponse = ResponseModel[
-    SessionCommand
+    ResponseTypes
 ]

robot-server/robot_server/service/session/session_types/base_session.py

@@ -64,19 +64,17 @@ async def clean_up(self):
         pass
 
     async def execute_command(self, command: command_models.RequestTypes) -> \
-            command_models.SessionCommand:
+            command_models.ResponseTypes:
         """Execute a command."""
         command_obj = create_command(command)
         command_result = await self.command_executor.execute(command_obj)
 
-        return command_models.SessionCommand(
-            id=command_obj.meta.identifier,
-            data=command.data,
-            command=command.command,
+        return command.make_response(
+            identifier=command_obj.meta.identifier,
             status=command_result.result.status,
-            createdAt=command_obj.meta.created_at,
-            startedAt=command_result.result.started_at,
-            completedAt=command_result.result.completed_at,
+            created_at=command_obj.meta.created_at,
+            started_at=command_result.result.started_at,
+            completed_at=command_result.result.completed_at,
             result=command_result.result.data,
         )