shared-data: deprecate json_encoder and use v2 serializer

Author: ahiuchingau

shared-data/python/opentrons_shared_data/labware/models.py

@@ -42,17 +42,13 @@ def __add__(self, other: Any) -> OffsetVector:
         """Adds two vectors together."""
         if not isinstance(other, OffsetVector):
             return NotImplemented
-        return OffsetVector(
-            x=self.x + other.x, y=self.y + other.y, z=self.z + other.z
-        )
+        return OffsetVector(x=self.x + other.x, y=self.y + other.y, z=self.z + other.z)
 
     def __sub__(self, other: Any) -> OffsetVector:
         """Subtracts two vectors."""
         if not isinstance(other, OffsetVector):
             return NotImplemented
-        return OffsetVector(
-            x=self.x - other.x, y=self.y - other.y, z=self.z - other.z
-        )
+        return OffsetVector(x=self.x - other.x, y=self.y - other.y, z=self.z - other.z)
 
 
 class GripperOffsets(BaseModel):

shared-data/python/opentrons_shared_data/pipette/pipette_definition.py

@@ -1,7 +1,13 @@
 import re
-from typing import List, Dict, Tuple, Optional, Annotated
-from pydantic import field_validator, ConfigDict, BaseModel, Field, BeforeValidator
-from typing_extensions import Literal
+from typing import List, Dict, Tuple, Optional, Annotated, Literal, TypeVar, Generic
+from pydantic import (
+    field_validator,
+    ConfigDict,
+    BaseModel,
+    Field,
+    BeforeValidator,
+    PlainSerializer,
+)
 from dataclasses import dataclass
 
 from . import types as pip_types, dev_types
@@ -21,6 +27,10 @@ def validate_opentrons_tiprack(v: str) -> str:
     return v
 
 
+EnumType = TypeVar("EnumType")
+EnumSerializer = Annotated[EnumType, PlainSerializer(lambda v: v.value)]
+
+
 # TODO (lc 12-5-2022) Ideally we can deprecate this
 # at somepoint once we load pipettes by channels and type
 @dataclass
@@ -286,12 +296,12 @@ class PipettePhysicalPropertiesDefinition(BaseModel):
         description="A list of pipette names that are compatible with this pipette.",
         alias="backCompatNames",
     )
-    pipette_type: pip_types.PipetteModelType = Field(
+    pipette_type: EnumSerializer[pip_types.PipetteModelType] = Field(
         ...,
         description="The pipette model type (related to number of channels).",
         alias="model",
     )
-    display_category: pip_types.PipetteGenerationType = Field(
+    display_category: EnumSerializer[pip_types.PipetteGenerationType] = Field(
         ..., description="The product model of the pipette.", alias="displayCategory"
     )
     pick_up_tip_configurations: PickUpTipConfigurations = Field(
@@ -313,7 +323,7 @@ class PipettePhysicalPropertiesDefinition(BaseModel):
     partial_tip_configurations: PartialTipDefinition = Field(
         ..., alias="partialTipConfigurations"
     )
-    channels: pip_types.PipetteChannelType = Field(
+    channels: EnumSerializer[pip_types.PipetteChannelType] = Field(
         ..., description="The maximum number of channels on the pipette."
     )
     shaft_diameter: float = Field(
@@ -329,7 +339,7 @@ class PipettePhysicalPropertiesDefinition(BaseModel):
         description="The distance of backlash on the plunger motor.",
         alias="backlashDistance",
     )
-    quirks: List[pip_types.Quirks] = Field(
+    quirks: List[EnumSerializer[pip_types.Quirks]] = Field(
         ..., description="The list of quirks available for the loaded configuration"
     )
     tip_presence_check_distance_mm: float = Field(
@@ -373,17 +383,6 @@ def convert_plunger_positions(
     ) -> Dict[pip_types.LiquidClasses, PlungerPositions]:
         return {pip_types.LiquidClasses[key]: value for key, value in v.items()}
 
-    # TODO[pydantic]: The following keys were removed: `json_encoders`.
-    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
-    model_config = ConfigDict(
-        json_encoders={
-            pip_types.PipetteChannelType: lambda v: v.value,
-            pip_types.PipetteModelType: lambda v: v.value,
-            pip_types.PipetteGenerationType: lambda v: v.value,
-            pip_types.Quirks: lambda v: v.value,
-        }
-    )
-
 
 class PipetteRowDefinition(BaseModel):
     key: str

shared-data/python/opentrons_shared_data/pipette/types.py

@@ -1,6 +1,6 @@
 import enum
 from dataclasses import dataclass
-from typing import Union, Dict, Mapping, Tuple, cast
+from typing import Union, Dict, Mapping, Tuple, TypeVar, cast
 from typing_extensions import Literal
 
 """Pipette Definition V2 Types"""

shared-data/python/opentrons_shared_data/protocol/models/protocol_schema_v6.py

@@ -1,4 +1,11 @@
-from pydantic import ConfigDict, BaseModel, Field, model_validator, field_validator, TypeAdapter
+from pydantic import (
+    ConfigDict,
+    BaseModel,
+    Field,
+    model_validator,
+    field_validator,
+    TypeAdapter,
+)
 from typing import Any, List, Optional, Dict, Union
 from typing_extensions import Literal
 from opentrons_shared_data.labware.models import LabwareDefinition

shared-data/python/tests/pipette/test_load_data.py

@@ -90,7 +90,7 @@ def test_update_pipette_configuration(
         base_configurations, v1_configuration_changes, liquid_class
     )
 
-    updated_configurations_dict = updated_configurations.dict()
+    updated_configurations_dict = updated_configurations.model_dump()
     for k, v in v1_configuration_changes.items():
         if k == "tip_length":
             for i in updated_configurations_dict["liquid_properties"][liquid_class][

shared-data/python/tests/pipette/test_validate_schema.py

@@ -1,11 +1,13 @@
 import os
+from typing import List
 from opentrons_shared_data import get_shared_data_root
 
 from opentrons_shared_data.pipette.pipette_definition import PipetteConfigurations
 from opentrons_shared_data.pipette.load_data import load_definition
 from opentrons_shared_data.pipette.pipette_load_name_conversions import (
     convert_pipette_model,
 )
+from opentrons_shared_data.pipette import types
 from opentrons_shared_data.pipette.dev_types import PipetteModel
 
 
@@ -68,3 +70,48 @@ def test_pick_up_configs_tip_count_keys() -> None:
                     == pick_up_tip_configs.speed_by_tip_count.keys()
                     == pick_up_tip_configs.current_by_tip_count.keys()
                 )
+
+
+def test_serializer() -> None:
+    """Verify that the serializer works as expected."""
+
+    loaded_model = load_definition(
+        types.PipetteModelType.p1000,
+        types.PipetteChannelType.NINETY_SIX_CHANNEL,
+        types.PipetteVersionType(3, 3),
+    )
+    quirk_0 = types.Quirks.pickupTipShake
+    quirk_1 = types.Quirks.dropTipShake
+    loaded_model.quirks = [quirk_0, quirk_1]
+
+    assert loaded_model.pipette_type == types.PipetteModelType.p1000
+    assert loaded_model.display_category == types.PipetteGenerationType.FLEX
+    assert loaded_model.channels == types.PipetteChannelType.NINETY_SIX_CHANNEL
+
+    model_dict = loaded_model.model_dump()
+    # each field should be the value of the enum class
+    assert (
+        isinstance(model_dict["pipette_type"], str)
+        and model_dict["pipette_type"] == loaded_model.pipette_type.value
+    )
+    assert (
+        isinstance(model_dict["display_category"], str)
+        and model_dict["display_category"] == loaded_model.display_category.value
+    )
+    assert (
+        isinstance(model_dict["channels"], int)
+        and model_dict["channels"] == loaded_model.channels.value
+    )
+
+    assert len(model_dict["quirks"]) == 2
+    dict_quirk_0 = model_dict["quirks"][0]
+    dict_quirk_1 = model_dict["quirks"][1]
+    assert isinstance(dict_quirk_0, str) and dict_quirk_0 == quirk_0.value
+    assert isinstance(dict_quirk_1, str) and dict_quirk_1 == quirk_1.value
+
+
+# TODO: (AA, 4/9/2024) we should add a test to validate the dumped json to make
+# sure we can re-load it as the BaseModel class. But we haven't added serializer
+# for other enums yet, such as LiquidClass, and since we haven't been
+# creating the definition files using model_dump/model_dump_json, it is okay to
+# skip this for now.