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,13 +38,19 @@ 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,
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))
     >>> Scalar(0)
-    nitypes.scalar.Scalar(value=0, units='')
+    nitypes.scalar.Scalar(value=0,
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))
     >>> Scalar(5.0, 'volts')
-    nitypes.scalar.Scalar(value=5.0, units='volts')
+    nitypes.scalar.Scalar(value=5.0,
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription':
+    'volts'}))
     >>> Scalar("value", "volts")
-    nitypes.scalar.Scalar(value='value', units='volts')
+    nitypes.scalar.Scalar(value='value',
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription':
+    'volts'}))
 
     Class members
     ^^^^^^^^^^^^^
@@ -60,12 +68,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 +91,15 @@ 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 the caller specifies a non-blank units string, overwrite the existing entry.
+        if UNIT_DESCRIPTION not in self._extended_properties or units:
+            self._extended_properties[UNIT_DESCRIPTION] = units
 
     @property
     def value(self) -> TScalar_co:
@@ -164,11 +185,23 @@ 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, self.units)
+        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}",
+            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,13 +36,19 @@ 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],
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))
     >>> Vector([0, 1, 2])
-    nitypes.vector.Vector(values=[0, 1, 2], units='')
+    nitypes.vector.Vector(values=[0, 1, 2],
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))
     >>> Vector([5.0, 6.0], 'volts')
-    nitypes.vector.Vector(values=[5.0, 6.0], units='volts')
+    nitypes.vector.Vector(values=[5.0, 6.0],
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription':
+    'volts'}))
     >>> Vector(["one", "two"], "volts")
-    nitypes.vector.Vector(values=['one', 'two'], units='volts')
+    nitypes.vector.Vector(values=['one', 'two'],
+    extended_properties=nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription':
+    'volts'}))
     """
 
     __slots__ = [
@@ -60,6 +67,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.
 
@@ -94,8 +103,15 @@ 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 the caller specifies a non-blank units string, overwrite the existing entry.
+        if UNIT_DESCRIPTION not in self._extended_properties or units:
+            self._extended_properties[UNIT_DESCRIPTION] = units
 
     @property
     def units(self) -> str:
@@ -191,11 +207,24 @@ 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, self.units)
+        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}",
+            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

@@ -182,14 +182,47 @@ 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, "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(False),
+            "nitypes.scalar.Scalar(value=False, extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Scalar(10),
+            "nitypes.scalar.Scalar(value=10, extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Scalar(20.0),
+            "nitypes.scalar.Scalar(value=20.0, extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Scalar("value"),
+            "nitypes.scalar.Scalar(value='value', extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Scalar(False, "amps"),
+            "nitypes.scalar.Scalar(value=False, extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': 'amps'}))",
+        ),
+        (
+            Scalar(10, "volts"),
+            "nitypes.scalar.Scalar(value=10, extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': 'volts'}))",
+        ),
+        (
+            Scalar(20.0, "watts"),
+            "nitypes.scalar.Scalar(value=20.0, extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': 'watts'}))",
+        ),
+        (
+            Scalar("value", ""),
+            "nitypes.scalar.Scalar(value='value', "
+            "extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
     ],
 )
 def test___various_values___repr___looks_ok(value: Scalar[Any], expected_repr: str) -> None:
@@ -244,12 +277,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 +298,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:

tests/unit/vector/test_vector.py

@@ -293,14 +293,46 @@ def test___different_units___comparison___not_equal() -> None:
 @pytest.mark.parametrize(
     "value, expected_repr",
     [
-        (Vector([False, True]), "nitypes.vector.Vector(values=[False, True], units='')"),
-        (Vector([10, 20]), "nitypes.vector.Vector(values=[10, 20], units='')"),
-        (Vector([20.0, 20.1]), "nitypes.vector.Vector(values=[20.0, 20.1], units='')"),
-        (Vector(["a", "b"]), "nitypes.vector.Vector(values=['a', 'b'], units='')"),
-        (Vector([False, True], "f"), "nitypes.vector.Vector(values=[False, True], units='f')"),
-        (Vector([10, 20], "volts"), "nitypes.vector.Vector(values=[10, 20], units='volts')"),
-        (Vector([20.0, 20.1], "w"), "nitypes.vector.Vector(values=[20.0, 20.1], units='w')"),
-        (Vector(["a", "b"], ""), "nitypes.vector.Vector(values=['a', 'b'], units='')"),
+        (
+            Vector([False, True]),
+            "nitypes.vector.Vector(values=[False, True], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Vector([10, 20]),
+            "nitypes.vector.Vector(values=[10, 20], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Vector([20.0, 20.1]),
+            "nitypes.vector.Vector(values=[20.0, 20.1], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Vector(["a", "b"]),
+            "nitypes.vector.Vector(values=['a', 'b'], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
+        (
+            Vector([False, True], "f"),
+            "nitypes.vector.Vector(values=[False, True], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': 'f'}))",
+        ),
+        (
+            Vector([10, 20], "volts"),
+            "nitypes.vector.Vector(values=[10, 20], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': 'volts'}))",
+        ),
+        (
+            Vector([20.0, 20.1], "w"),
+            "nitypes.vector.Vector(values=[20.0, 20.1], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': 'w'}))",
+        ),
+        (
+            Vector(["a", "b"], ""),
+            "nitypes.vector.Vector(values=['a', 'b'], extended_properties="
+            "nitypes.waveform.ExtendedPropertyDictionary({'NI_UnitDescription': ''}))",
+        ),
     ],
 )
 def test___various_values___repr___looks_ok(value: Vector[Any], expected_repr: str) -> None:
@@ -359,12 +391,14 @@ def test___vector_with_units___set_units___units_updated_correctly() -> None:
         Vector([10, 20], "volts"),
         Vector([20.0, 20.1], "watts"),
         Vector(["a", "b"], ""),
+        Vector([10, 20], "volts", extended_properties={"one": 1}),
     ],
 )
 def test___various_values___copy___makes_copy(value: Vector[TScalar]) -> 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(
@@ -378,12 +412,15 @@ def test___various_values___copy___makes_copy(value: Vector[TScalar]) -> None:
         Vector([10, 20], "volts"),
         Vector([20.0, 20.1], "watts"),
         Vector(["a", "b"], ""),
+        Vector([10, 20], "volts", extended_properties={"one": 1}),
     ],
 )
 def test___various_values___pickle_unpickle___makes_copy(value: Vector[TScalar]) -> None:
     new_value = pickle.loads(pickle.dumps(value))
+    assert isinstance(new_value, Vector)
     assert new_value is not value
     assert new_value == value
+    assert new_value.extended_properties == value.extended_properties
 
 
 def test___vector___pickle___references_public_modules() -> None: