diff --git a/api/src/opentrons/protocols/api_support/util.py b/api/src/opentrons/protocols/api_support/util.py index 70adacad775..5ec5cc93651 100644 --- a/api/src/opentrons/protocols/api_support/util.py +++ b/api/src/opentrons/protocols/api_support/util.py @@ -16,6 +16,8 @@ TypeVar, Union, cast, + KeysView, + ItemsView, ) from opentrons import types as top_types @@ -289,7 +291,7 @@ def __getitem__(self, key: Union[str, Axis]) -> float: return self.data[checked_key] @staticmethod - def _verify_key(key: Any) -> Axis: + def _verify_key(key: object) -> Axis: if isinstance(key, Axis): checked_key: Optional[Axis] = key elif isinstance(key, str): @@ -300,56 +302,40 @@ def _verify_key(key: Any) -> Axis: raise KeyError(key) return checked_key - def __setitem__(self, key: Any, value: Any) -> None: + def __setitem__(self, key: object, value: object) -> None: + + checked_key = AxisMaxSpeeds._verify_key(key) if value is None: - del self[key] + del self[checked_key] return - checked_key = AxisMaxSpeeds._verify_key(key) checked_val = _assert_gzero( value, "max speeds should be numerical values in mm/s" ) self.data[checked_key] = checked_val + def _axis_to_string(self, axis: Union[str, Axis]) -> str: + if isinstance(axis, str): + return axis + if self._robot_type == "OT-3 Standard": + return axis.name + return ot2_axis_to_string(axis) + def __delitem__(self, key: Union[str, Axis]) -> None: checked_key = AxisMaxSpeeds._verify_key(key) del self.data[checked_key] def __iter__(self) -> Iterator[str]: """keys() and dict iteration return string keys""" - string_keys = ( - k - if isinstance(k, str) - else k.name - if self._robot_type == "OT-3 Standard" - else ot2_axis_to_string(k) - for k in self.data.keys() - ) + string_keys = (self._axis_to_string(k) for k in self.data.keys()) return string_keys - def keys(self) -> Iterator[str | Axis]: # type: ignore[override] - string_keys = ( - k - if isinstance(k, str) - else k.name - if self._robot_type == "OT-3 Standard" - else ot2_axis_to_string(k) - for k in self.data.keys() - ) - return string_keys + def keys(self) -> KeysView[str]: + return ({self._axis_to_string(k): v for k, v in self.data.items()}).keys() - def items(self) -> Iterator[Dict[str | Axis, float]]: # type: ignore[override] - return ( - { - k - if isinstance(k, str) - else k.name - if self._robot_type == "OT-3 Standard" - else ot2_axis_to_string(k): v - } - for k, v in self.data.items() - ) + def items(self) -> ItemsView[str, float]: + return ({self._axis_to_string(k): v for k, v in self.data.items()}).items() def clamp_value( diff --git a/api/tests/opentrons/protocols/api_support/test_util.py b/api/tests/opentrons/protocols/api_support/test_util.py index 87de52e76eb..7335722e3e2 100644 --- a/api/tests/opentrons/protocols/api_support/test_util.py +++ b/api/tests/opentrons/protocols/api_support/test_util.py @@ -47,7 +47,7 @@ def test_max_speeds_userdict() -> None: assert defaults["a"] == 20 assert defaults[Axis.A] == 20 - assert sorted(list(defaults.keys())) == sorted(["X", "A"]) # type: ignore[type-var] + assert sorted(list(defaults.keys())) == sorted(["X", "A"]) assert "X" in defaults.keys() del defaults["A"]