Enforce typing and type checking

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

@@ -1,11 +1,20 @@
 """Helper functions for liquid-level related calculations inside a given frustum."""
-from typing import List, Tuple, Any, Optional
+from typing import List, Tuple, Optional, Union
 from numpy import pi, iscomplex, roots, real
 from math import isclose
 
 from ..errors.exceptions import InvalidLiquidHeightFound, InvalidWellDefinitionError
 
-from opentrons_shared_data.labware.labware_definition import InnerWellGeometry
+from opentrons_shared_data.labware.types import (
+    InnerWellGeometry,
+    WellSegment,
+    ToWellSegmentDict,
+    ToInnerWellGeometryDict,
+)
+from opentrons_shared_data.labware.labware_definition import (
+    InnerWellGeometry as InnerWellGeometryDef,
+    WellSegment as WellSegmentDef,
+)
 
 
 def _reject_unacceptable_heights(
@@ -206,15 +215,18 @@ def _height_from_volume_spherical(
 
 
 def get_well_volumetric_capacity(
-    well_geometry: InnerWellGeometry,
+    well_geometry: Union[InnerWellGeometry, InnerWellGeometryDef],
 ) -> List[Tuple[float, float]]:
     """Return the total volumetric capacity of a well as a map of height borders to volume."""
+    checked_geometry = ToInnerWellGeometryDict(well_geometry)
     # dictionary map of heights to volumetric capacities within their respective segment
     # {top_height_0: volume_0, top_height_1: volume_1, top_height_2: volume_2}
     well_volume = []
 
     # get the well segments sorted in ascending order
-    sorted_well = sorted(well_geometry.sections, key=lambda section: section.topHeight)
+    sorted_well = sorted(
+        checked_geometry["sections"], key=lambda section: section["topHeight"]
+    )
 
     for segment in sorted_well:
         section_volume: Optional[float] = None
@@ -255,32 +267,33 @@ def get_well_volumetric_capacity(
 
 
 def height_at_volume_within_section(
-    section: Any,
+    section: Union[WellSegment, WellSegmentDef],
     target_volume_relative: float,
     section_height: float,
 ) -> float:
     """Calculate a height within a bounded section according to geometry."""
-    if section["shape"] == "spherical":
+    checked_section = ToWellSegmentDict(section)
+    if checked_section["shape"] == "spherical":
         partial_height = _height_from_volume_spherical(
             volume=target_volume_relative,
             total_frustum_height=section_height,
-            radius_of_curvature=section.radiusOfCurvature,
+            radius_of_curvature=checked_section["radiusOfCurvature"],
         )
-    elif section["shape"] == "circular":
+    elif checked_section["shape"] == "circular":
         partial_height = _height_from_volume_circular(
             volume=target_volume_relative,
-            top_radius=(section["bottomDiameter"] / 2),
-            bottom_radius=(section["topDiameter"] / 2),
+            top_radius=(checked_section["bottomDiameter"] / 2),
+            bottom_radius=(checked_section["topDiameter"] / 2),
             total_frustum_height=section_height,
         )
-    elif section["shape"] == "rectangular":
+    elif checked_section["shape"] == "rectangular":
         partial_height = _height_from_volume_rectangular(
             volume=target_volume_relative,
             total_frustum_height=section_height,
-            bottom_width=section["bottomXDimension"],
-            bottom_length=section["bottomYDimension"],
-            top_width=section["topXDimension"],
-            top_length=section["topYDimension"],
+            bottom_width=checked_section["bottomXDimension"],
+            bottom_length=checked_section["bottomYDimension"],
+            top_width=checked_section["topXDimension"],
+            top_length=checked_section["topYDimension"],
         )
     else:
         raise NotImplementedError(
@@ -290,31 +303,32 @@ def height_at_volume_within_section(
 
 
 def volume_at_height_within_section(
-    section: Any,
+    section: Union[WellSegment, WellSegmentDef],
     target_height_relative: float,
     section_height: float,
 ) -> float:
     """Calculate a volume within a bounded section according to geometry."""
-    if section["shape"] == "spherical":
+    checked_section = ToWellSegmentDict(section)
+    if checked_section["shape"] == "spherical":
         partial_volume = _volume_from_height_spherical(
             target_height=target_height_relative,
-            radius_of_curvature=section["radiusOfCurvature"],
+            radius_of_curvature=checked_section["radiusOfCurvature"],
         )
-    elif section["shape"] == "circular":
+    elif checked_section["shape"] == "circular":
         partial_volume = _volume_from_height_circular(
             target_height=target_height_relative,
             total_frustum_height=section_height,
-            bottom_radius=(section["bottomDiameter"] / 2),
-            top_radius=(section["topDiameter"] / 2),
+            bottom_radius=(checked_section["bottomDiameter"] / 2),
+            top_radius=(checked_section["topDiameter"] / 2),
         )
-    elif section["shape"] == "rectangular":
+    elif checked_section["shape"] == "rectangular":
         partial_volume = _volume_from_height_rectangular(
             target_height=target_height_relative,
             total_frustum_height=section_height,
-            bottom_width=section["BottomXDimension"],
-            bottom_length=section["BottomYDimension"],
-            top_width=section["TopXDimension"],
-            top_length=section["TopYDimension"],
+            bottom_width=checked_section["bottomXDimension"],
+            bottom_length=checked_section["bottomYDimension"],
+            top_width=checked_section["topXDimension"],
+            top_length=checked_section["topYDimension"],
         )
     # TODO(cm): this would be the NEST-96 2uL wells referenced in EXEC-712
     # we need to input the math attached to that issue
@@ -326,7 +340,7 @@ def volume_at_height_within_section(
 
 
 def _find_volume_in_partial_frustum(
-    sorted_well: List[Any],
+    sorted_well: List[WellSegment],
     target_height: float,
 ) -> Optional[float]:
     """Look through a sorted list of frusta for a target height, and find the volume at that height."""
@@ -349,10 +363,11 @@ def _find_volume_in_partial_frustum(
 
 
 def find_volume_at_well_height(
-    target_height: float, well_geometry: InnerWellGeometry
+    target_height: float, well_geometry: Union[InnerWellGeometry, InnerWellGeometryDef]
 ) -> float:
     """Find the volume within a well, at a known height."""
-    volumetric_capacity = get_well_volumetric_capacity(well_geometry)
+    checked_geometry = ToInnerWellGeometryDict(well_geometry)
+    volumetric_capacity = get_well_volumetric_capacity(checked_geometry)
     max_height = volumetric_capacity[-1][0]
     if target_height < 0 or target_height > max_height:
         raise InvalidLiquidHeightFound("Invalid target height.")
@@ -370,7 +385,9 @@ def find_volume_at_well_height(
     # find the section the target height is in and compute the volume
     # since bottomShape is not in list of frusta, check here first
 
-    sorted_well = sorted(well_geometry.sections, key=lambda section: section.topHeight)
+    sorted_well = sorted(
+        checked_geometry["sections"], key=lambda section: section["topHeight"]
+    )
     # TODO(cm): handle non-frustum section that is not at the bottom.
     partial_volume = _find_volume_in_partial_frustum(
         sorted_well=sorted_well,
@@ -382,7 +399,7 @@ def find_volume_at_well_height(
 
 
 def _find_height_in_partial_frustum(
-    sorted_well: List[Any],
+    sorted_well: List[WellSegment],
     volumetric_capacity: List[Tuple[float, float]],
     target_volume: float,
 ) -> Optional[float]:
@@ -407,18 +424,21 @@ def _find_height_in_partial_frustum(
 
 
 def find_height_at_well_volume(
-    target_volume: float, well_geometry: InnerWellGeometry
+    target_volume: float, well_geometry: Union[InnerWellGeometry, InnerWellGeometryDef]
 ) -> float:
     """Find the height within a well, at a known volume."""
-    volumetric_capacity = get_well_volumetric_capacity(well_geometry)
+    checked_geometry = ToInnerWellGeometryDict(well_geometry)
+    volumetric_capacity = get_well_volumetric_capacity(checked_geometry)
     max_volume = volumetric_capacity[-1][1]
     if target_volume < 0 or target_volume > max_volume:
         raise InvalidLiquidHeightFound("Invalid target volume.")
     for section_height, section_volume in volumetric_capacity:
         if target_volume == section_volume:
             return section_height
 
-    sorted_well = sorted(well_geometry.sections, key=lambda section: section.topHeight)
+    sorted_well = sorted(
+        checked_geometry["sections"], key=lambda section: section["topHeight"]
+    )
     # find the section the target volume is in and compute the height
     well_height = _find_height_in_partial_frustum(
         sorted_well=sorted_well,