Fix copy/pickling for Scalar and Vector

src/nitypes/scalar.py

@@ -9,12 +9,14 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Generic, Union
 
-from typing_extensions import TypeVar, final
+from typing_extensions import Self, TypeVar, final
 
 from nitypes._exceptions import invalid_arg_type
 from nitypes.waveform._extended_properties import UNIT_DESCRIPTION
+from nitypes.waveform.typing import ExtendedPropertyValue
 
 if TYPE_CHECKING:
     # Import from the public package so the docs don't reference private submodules.
@@ -36,9 +38,9 @@ class Scalar(Generic[TScalar_co]):
     To construct a scalar data object, use the :class:`Scalar` class:
 
     >>> Scalar(False)
-    nitypes.scalar.Scalar(value=False, units='')
+    nitypes.scalar.Scalar(value=False)
     >>> Scalar(0)
-    nitypes.scalar.Scalar(value=0, units='')
+    nitypes.scalar.Scalar(value=0)
     >>> Scalar(5.0, 'volts')
     nitypes.scalar.Scalar(value=5.0, units='volts')
     >>> Scalar("value", "volts")
@@ -60,12 +62,18 @@ def __init__(
         self,
         value: TScalar_co,
         units: str = "",
+        *,
+        extended_properties: Mapping[str, ExtendedPropertyValue] | None = None,
+        copy_extended_properties: bool = True,
     ) -> None:
         """Initialize a new scalar.
 
         Args:
             value: The scalar data to store in this object.
             units: The units string associated with this data.
+            extended_properties: The extended properties of the Scalar.
+            copy_extended_properties: Specifies whether to copy the extended properties or take
+                ownership.
 
         Returns:
             A scalar data object.
@@ -77,8 +85,20 @@ def __init__(
             raise invalid_arg_type("units", "str", units)
 
         self._value = value
-        self._extended_properties = ExtendedPropertyDictionary()
-        self._extended_properties[UNIT_DESCRIPTION] = units
+        if copy_extended_properties or not isinstance(
+            extended_properties, ExtendedPropertyDictionary
+        ):
+            extended_properties = ExtendedPropertyDictionary(extended_properties)
+        self._extended_properties = extended_properties
+
+        # If units are not already in extended properties, set them.
+        if UNIT_DESCRIPTION not in self._extended_properties:
+            self._extended_properties[UNIT_DESCRIPTION] = units
+        elif units and units != self._extended_properties.get(UNIT_DESCRIPTION):
+            raise ValueError(
+                "The specified units input does not match the units specified in "
+                "extended_properties."
+            )
 
     @property
     def value(self) -> TScalar_co:
@@ -164,11 +184,28 @@ def __le__(self, value: Scalar[TScalar_co]) -> bool:
 
     def __reduce__(self) -> tuple[Any, ...]:
         """Return object state for pickling."""
-        return (self.__class__, (self.value, self.units))
+        ctor_args = (self.value,)
+        ctor_kwargs: dict[str, Any] = {
+            "extended_properties": self._extended_properties,
+            "copy_extended_properties": False,
+        }
+        return (self.__class__._unpickle, (ctor_args, ctor_kwargs))
+
+    @classmethod
+    def _unpickle(cls, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Self:
+        return cls(*args, **kwargs)
 
     def __repr__(self) -> str:
         """Return repr(self)."""
-        args = [f"value={self.value!r}", f"units={self.units!r}"]
+        args = [f"value={self.value!r}"]
+
+        if self.units:
+            args.append(f"units={self.units!r}")
+
+        # Only display the extended properties if non-units entries are specified.
+        if any(key for key in self.extended_properties.keys() if key != UNIT_DESCRIPTION):
+            args.append(f"extended_properties={self.extended_properties!r}")
+
         return f"{self.__class__.__module__}.{self.__class__.__name__}({', '.join(args)})"
 
     def __str__(self) -> str:

src/nitypes/vector.py

@@ -8,13 +8,14 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, MutableSequence
-from typing import TYPE_CHECKING, overload, Any, Union
+from collections.abc import Iterable, Mapping, MutableSequence
+from typing import TYPE_CHECKING, Any, Union, overload
 
-from typing_extensions import TypeVar, final, override
+from typing_extensions import Self, TypeVar, final, override
 
 from nitypes._exceptions import invalid_arg_type
 from nitypes.waveform._extended_properties import UNIT_DESCRIPTION
+from nitypes.waveform.typing import ExtendedPropertyValue
 
 if TYPE_CHECKING:
     # Import from the public package so the docs don't reference private submodules.
@@ -35,9 +36,9 @@ class Vector(MutableSequence[TScalar]):
     To construct a vector data object, use the :class:`Vector` class:
 
     >>> Vector([False, True])
-    nitypes.vector.Vector(values=[False, True], units='')
+    nitypes.vector.Vector(values=[False, True])
     >>> Vector([0, 1, 2])
-    nitypes.vector.Vector(values=[0, 1, 2], units='')
+    nitypes.vector.Vector(values=[0, 1, 2])
     >>> Vector([5.0, 6.0], 'volts')
     nitypes.vector.Vector(values=[5.0, 6.0], units='volts')
     >>> Vector(["one", "two"], "volts")
@@ -60,6 +61,8 @@ def __init__(
         units: str = "",
         *,
         value_type: type[TScalar] | None = None,
+        extended_properties: Mapping[str, ExtendedPropertyValue] | None = None,
+        copy_extended_properties: bool = True,
     ) -> None:
         """Initialize a new vector.
 
@@ -69,6 +72,9 @@ def __init__(
             value_type: The type of values that will be added to this Vector.
                 This parameter should only be used when creating a Vector with
                 an empty Iterable.
+            extended_properties: The extended properties of the Vector.
+            copy_extended_properties: Specifies whether to copy the extended properties or take
+                ownership.
 
         Returns:
             A vector data object.
@@ -94,8 +100,20 @@ def __init__(
             raise invalid_arg_type("units", "str", units)
 
         self._values = list(values)
-        self._extended_properties = ExtendedPropertyDictionary()
-        self._extended_properties[UNIT_DESCRIPTION] = units
+        if copy_extended_properties or not isinstance(
+            extended_properties, ExtendedPropertyDictionary
+        ):
+            extended_properties = ExtendedPropertyDictionary(extended_properties)
+        self._extended_properties = extended_properties
+
+        # If units are not already in extended properties, set them.
+        if UNIT_DESCRIPTION not in self._extended_properties:
+            self._extended_properties[UNIT_DESCRIPTION] = units
+        elif units and units != self._extended_properties.get(UNIT_DESCRIPTION):
+            raise ValueError(
+                "The specified units input does not match the units specified in "
+                "extended_properties."
+            )
 
     @property
     def units(self) -> str:
@@ -191,11 +209,29 @@ def __eq__(self, value: object, /) -> bool:
 
     def __reduce__(self) -> tuple[Any, ...]:
         """Return object state for pickling."""
-        return (self.__class__, (self._values, self.units))
+        ctor_args = (self._values,)
+        ctor_kwargs: dict[str, Any] = {
+            "value_type": self._value_type,
+            "extended_properties": self._extended_properties,
+            "copy_extended_properties": False,
+        }
+        return (self.__class__._unpickle, (ctor_args, ctor_kwargs))
+
+    @classmethod
+    def _unpickle(cls, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Self:
+        return cls(*args, **kwargs)
 
     def __repr__(self) -> str:
         """Return repr(self)."""
-        args = [f"values={self._values!r}", f"units={self.units!r}"]
+        args = [f"values={self._values!r}"]
+
+        if self.units:
+            args.append(f"units={self.units!r}")
+
+        # Only display the extended properties if non-units entries are specified.
+        if any(key for key in self.extended_properties.keys() if key != UNIT_DESCRIPTION):
+            args.append(f"extended_properties={self.extended_properties!r}")
+
         return f"{self.__class__.__module__}.{self.__class__.__name__}({', '.join(args)})"
 
     def __str__(self) -> str:

src/nitypes/xy_data.py

@@ -71,20 +71,16 @@ class XYData(Generic[TData]):
     To construct an XYData object, use the :class:`XYData` class:
 
     >>> XYData(np.array([1.1], np.float64), np.array([4.1], np.float64))
-    nitypes.xy_data.XYData(x_data=array([1.1]), y_data=array([4.1]),
-    extended_properties={'NI_UnitDescription_X': '', 'NI_UnitDescription_Y': ''})
+    nitypes.xy_data.XYData(x_data=array([1.1]), y_data=array([4.1]))
     >>> XYData(np.array([1, 2]), np.array([4, 5]), x_units="A", y_units="V")
-    nitypes.xy_data.XYData(x_data=array([1, 2]), y_data=array([4, 5]),
-    extended_properties={'NI_UnitDescription_X': 'A', 'NI_UnitDescription_Y': 'V'})
+    nitypes.xy_data.XYData(x_data=array([1, 2]), y_data=array([4, 5]), x_units='A', y_units='V')
 
     To construct an XYData object using built-in lists, use from_arrays_1d():
 
     >>> XYData.from_arrays_1d([1, 2], [5, 6], np.int32)
-    nitypes.xy_data.XYData(x_data=array([1, 2], dtype=int32), y_data=array([5, 6], dtype=int32),
-    extended_properties={'NI_UnitDescription_X': '', 'NI_UnitDescription_Y': ''})
+    nitypes.xy_data.XYData(x_data=array([1, 2], dtype=int32), y_data=array([5, 6], dtype=int32))
     >>> XYData.from_arrays_1d([1.0, 1.1], [1.2, 1.3], np.float64)
-    nitypes.xy_data.XYData(x_data=array([1. , 1.1]), y_data=array([1.2, 1.3]),
-    extended_properties={'NI_UnitDescription_X': '', 'NI_UnitDescription_Y': ''})
+    nitypes.xy_data.XYData(x_data=array([1. , 1.1]), y_data=array([1.2, 1.3]))
     """
 
     __slots__ = [
@@ -246,12 +242,23 @@ def __init__(
             extended_properties = ExtendedPropertyDictionary(extended_properties)
         self._extended_properties = extended_properties
 
-        # If x and y units are not already in extended properties, set them.
-        # If the caller specifies a non-blank x or y units, overwrite the existing entry.
-        if _UNIT_DESCRIPTION_X not in self._extended_properties or x_units:
+        # If x_units are not already in extended properties, set them.
+        if _UNIT_DESCRIPTION_X not in self._extended_properties:
             self._extended_properties[_UNIT_DESCRIPTION_X] = x_units
-        if _UNIT_DESCRIPTION_Y not in self._extended_properties or y_units:
+        elif x_units and x_units != self._extended_properties.get(_UNIT_DESCRIPTION_X):
+            raise ValueError(
+                "The specified x_units input does not match the units specified in "
+                "extended_properties."
+            )
+
+        # If y_units are not already in extended properties, set them.
+        if _UNIT_DESCRIPTION_Y not in self._extended_properties:
             self._extended_properties[_UNIT_DESCRIPTION_Y] = y_units
+        elif y_units and y_units != self._extended_properties.get(_UNIT_DESCRIPTION_Y):
+            raise ValueError(
+                "The specified y_units input does not match the units specified in "
+                "extended_properties."
+            )
 
     def _init_with_provided_arrays(
         self,
@@ -365,11 +372,22 @@ def _unpickle(cls, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Self:
 
     def __repr__(self) -> str:
         """Return repr(self)."""
-        args = [
-            f"x_data={self.x_data!r}",
-            f"y_data={self.y_data!r}",
-            f"extended_properties={self._extended_properties._properties!r}",
-        ]
+        args = [f"x_data={self.x_data!r}", f"y_data={self.y_data!r}"]
+
+        if self.x_units:
+            args.append(f"x_units={self.x_units!r}")
+
+        if self.y_units:
+            args.append(f"y_units={self.y_units!r}")
+
+        # Only display the extended properties if non-units entries are specified.
+        if any(
+            key
+            for key in self.extended_properties.keys()
+            if key not in [_UNIT_DESCRIPTION_X, _UNIT_DESCRIPTION_Y]
+        ):
+            args.append(f"extended_properties={self.extended_properties!r}")
+
         return f"{self.__class__.__module__}.{self.__class__.__name__}({', '.join(args)})"
 
     def __str__(self) -> str:

tests/unit/scalar/test_scalar.py

@@ -68,6 +68,20 @@ def test___invalid_data_value___create___raises_type_error(data_value: Any) -> N
     assert exc.value.args[0].startswith("The scalar input data must be a bool, int, float, or str.")
 
 
+def test___both_units_specified_unequal__create___raises_value_error() -> None:
+    with pytest.raises(ValueError) as exc:
+        _ = Scalar(10, "Volts", extended_properties={UNIT_DESCRIPTION: "Amps"})
+
+    assert exc.value.args[0].startswith(
+        "The specified units input does not match the units specified in extended_properties."
+    )
+
+
+def test___units_only_specified_in_extended_properties__create___creates_with_units() -> None:
+    data = Scalar(10, extended_properties={UNIT_DESCRIPTION: "Volts"})
+    assert data.units == "Volts"
+
+
 ###############################################################################
 # compare
 ###############################################################################
@@ -182,14 +196,20 @@ def test___different_units___comparison___throws_exception() -> None:
 @pytest.mark.parametrize(
     "value, expected_repr",
     [
-        (Scalar(False), "nitypes.scalar.Scalar(value=False, units='')"),
-        (Scalar(10), "nitypes.scalar.Scalar(value=10, units='')"),
-        (Scalar(20.0), "nitypes.scalar.Scalar(value=20.0, units='')"),
-        (Scalar("value"), "nitypes.scalar.Scalar(value='value', units='')"),
+        (Scalar(False), "nitypes.scalar.Scalar(value=False)"),
+        (Scalar(10), "nitypes.scalar.Scalar(value=10)"),
+        (Scalar(20.0), "nitypes.scalar.Scalar(value=20.0)"),
+        (Scalar("value"), "nitypes.scalar.Scalar(value='value')"),
         (Scalar(False, "amps"), "nitypes.scalar.Scalar(value=False, units='amps')"),
         (Scalar(10, "volts"), "nitypes.scalar.Scalar(value=10, units='volts')"),
         (Scalar(20.0, "watts"), "nitypes.scalar.Scalar(value=20.0, units='watts')"),
-        (Scalar("value", ""), "nitypes.scalar.Scalar(value='value', units='')"),
+        (Scalar("value", ""), "nitypes.scalar.Scalar(value='value')"),
+        (
+            Scalar(10, units="volts", extended_properties={"Prop1": "Value1"}),
+            "nitypes.scalar.Scalar(value=10, units='volts', "
+            "extended_properties=nitypes.waveform.ExtendedPropertyDictionary("
+            "{'Prop1': 'Value1', 'NI_UnitDescription': 'volts'}))",
+        ),
     ],
 )
 def test___various_values___repr___looks_ok(value: Scalar[Any], expected_repr: str) -> None:
@@ -244,12 +264,14 @@ def test___scalar_with_units___set_units___units_updated_correctly() -> None:
         Scalar(10, "volts"),
         Scalar(20.0, "watts"),
         Scalar("value", ""),
+        Scalar(10, "Volts", extended_properties={"one": 1}),
     ],
 )
 def test___various_values___copy___makes_copy(value: Scalar[TScalar_co]) -> None:
     new_value = copy.copy(value)
     assert new_value is not value
     assert new_value == value
+    assert new_value.extended_properties == value.extended_properties
 
 
 @pytest.mark.parametrize(
@@ -263,12 +285,15 @@ def test___various_values___copy___makes_copy(value: Scalar[TScalar_co]) -> None
         Scalar(10, "volts"),
         Scalar(20.0, "watts"),
         Scalar("value", ""),
+        Scalar(10, "Volts", extended_properties={"one": 1}),
     ],
 )
 def test___various_values___pickle_unpickle___makes_copy(value: Scalar[TScalar_co]) -> None:
     new_value = pickle.loads(pickle.dumps(value))
+    assert isinstance(new_value, Scalar)
     assert new_value is not value
     assert new_value == value
+    assert new_value.extended_properties == value.extended_properties
 
 
 def test___scalar___pickle___references_public_modules() -> None: