waveform: Implement copy and pickle (#13)

* tests: Add copy/pickle test cases

* pyproject.toml: Use dev branch of hightime for copy/pickle fixes

* waveform: Fix waveform timing copy/pickle

* tests: Add more scale mode copy/pickle tests

* waveform: Fix AnalogWaveform copy/pickle

* waveform: Use Self instead of duplicating type

* waveform: Document that waveform timing objects are immutable and copy the timestamps list by default

* tests: Refactor some duplicated params

* waveform: Use _typing.Self

* waveform: Enable shallow copy of extended properties

* Only use dev branch of hightime for testing, not in distribution package

* pyproject.toml: Use hightime main branch for testing

* Update poetry.lock

* tests: Use assert_deep_copy in pickle test

Author: bkeryan

poetry.lock

@@ -14,8 +14,7 @@ numpy = [
   { version = ">=1.26", python = ">=3.12,<3.13" },
   { version = ">=2.1", python = "^3.13" },
 ]
-# hightime = "^0.2.2"
-hightime = { git = "https://github.com/ni/hightime.git" }
+hightime = "^0.2.2"
 
 [tool.poetry.group.lint.dependencies]
 bandit = { version = ">=1.7", extras = ["toml"] }
@@ -26,6 +25,8 @@ mypy = ">=1.0"
 pytest = ">=7.2"
 pytest-cov = ">=4.0"
 pytest-mock = ">=3.0"
+# Use an unreleased version of hightime for testing.
+hightime = { git = "https://github.com/ni/hightime.git" }
 
 [tool.poetry.group.docs]
 optional = true

pyproject.toml

@@ -12,7 +12,7 @@
 
 from nitypes._arguments import arg_to_uint, validate_dtype, validate_unsupported_arg
 from nitypes._exceptions import invalid_arg_type, invalid_array_ndim
-from nitypes._typing import TypeAlias
+from nitypes._typing import Self, TypeAlias
 from nitypes.waveform._extended_properties import (
     CHANNEL_NAME,
     UNIT_DESCRIPTION,
@@ -351,6 +351,7 @@ def __init__(
         start_index: SupportsIndex | None = None,
         capacity: SupportsIndex | None = None,
         extended_properties: Mapping[str, ExtendedPropertyValue] | None = None,
+        copy_extended_properties: bool = True,
         timing: Timing | PrecisionTiming | None = None,
         scale_mode: ScaleMode | None = None,
     ) -> None:
@@ -368,6 +369,8 @@ def __init__(
             capacity: The number of samples to allocate. Pre-allocating a larger buffer optimizes
                 appending samples to the waveform.
             extended_properties: The extended properties of the analog waveform.
+            copy_extended_properties: Specifies whether to copy the extended properties or take
+                ownership.
             timing: The timing information of the analog waveform.
             scale_mode: The scale mode of the analog waveform.
 
@@ -389,7 +392,11 @@ def __init__(
         else:
             raise invalid_arg_type("raw data", "NumPy ndarray", raw_data)
 
-        self._extended_properties = ExtendedPropertyDictionary(extended_properties)
+        if copy_extended_properties or not isinstance(
+            extended_properties, ExtendedPropertyDictionary
+        ):
+            extended_properties = ExtendedPropertyDictionary(extended_properties)
+        self._extended_properties = extended_properties
 
         if timing is None:
             timing = Timing.empty
@@ -864,6 +871,22 @@ def __eq__(self, value: object, /) -> bool:
             and self._scale_mode == value._scale_mode
         )
 
+    def __reduce__(self) -> tuple[Any, ...]:
+        """Return object state for pickling."""
+        ctor_args = (self._sample_count, self.dtype)
+        ctor_kwargs: dict[str, Any] = {
+            "raw_data": self.raw_data,
+            "extended_properties": self._extended_properties,
+            "copy_extended_properties": False,
+            "timing": self._timing,
+            "scale_mode": self._scale_mode,
+        }
+        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"{self._sample_count}"]

src/nitypes/waveform/_analog_waveform.py

@@ -19,6 +19,8 @@
 class ExtendedPropertyDictionary(MutableMapping[str, ExtendedPropertyValue]):
     """A dictionary of extended properties."""
 
+    __slots__ = ["_properties"]
+
     def __init__(self, properties: Mapping[str, ExtendedPropertyValue] | None = None, /) -> None:
         """Construct an ExtendedPropertyDictionary."""
         self._properties: dict[str, ExtendedPropertyValue] = {}

src/nitypes/waveform/_extended_properties.py

@@ -4,7 +4,7 @@
 import operator
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
-from typing import Generic, SupportsIndex, TypeVar
+from typing import Any, Generic, SupportsIndex, TypeVar
 
 from nitypes._exceptions import add_note
 from nitypes._typing import Self
@@ -14,13 +14,15 @@
     create_sample_interval_strategy,
 )
 
-
 _TDateTime = TypeVar("_TDateTime", bound=dt.datetime)
 _TTimeDelta = TypeVar("_TTimeDelta", bound=dt.timedelta)
 
 
 class BaseTiming(ABC, Generic[_TDateTime, _TTimeDelta]):
-    """Base class for waveform timing information."""
+    """Base class for waveform timing information.
+
+    Waveform timing objects are immutable.
+    """
 
     @classmethod
     @abstractmethod
@@ -118,8 +120,26 @@ def __init__(
         time_offset: _TTimeDelta | None,
         sample_interval: _TTimeDelta | None,
         timestamps: Sequence[_TDateTime] | None,
+        *,
+        copy_timestamps: bool = True,
     ) -> None:
-        """Construct a base waveform timing object."""
+        """Construct a waveform timing object.
+
+        Args:
+            sample_interval_mode: The sample interval mode of the waveform timing.
+            timestamp: The timestamp of the waveform timing. This argument is optional for
+                SampleIntervalMode.NONE and SampleIntervalMode.REGULAR and unsupported for
+                SampleIntervalMode.IRREGULAR.
+            time_offset: The time difference between the timestamp and the first sample. This
+                argument is optional for SampleIntervalMode.NONE and SampleIntervalMode.REGULAR and
+                unsupported for SampleIntervalMode.IRREGULAR.
+            sample_interval: The time interval between samples. This argument is required for
+                SampleIntervalMode.REGULAR and unsupported otherwise.
+            timestamps: A sequence containing a timestamp for each sample in the waveform,
+                specifying the time that the sample was acquired. This argument is required for
+                SampleIntervalMode.IRREGULAR and unsupported otherwise.
+            copy_timestamps: Specifies whether to copy the timestamps or take ownership.
+        """
         sample_interval_strategy = create_sample_interval_strategy(sample_interval_mode)
         try:
             sample_interval_strategy.validate_init_args(
@@ -129,7 +149,7 @@ def __init__(
             add_note(e, f"Sample interval mode: {sample_interval_mode}")
             raise
 
-        if timestamps is not None and not isinstance(timestamps, list):
+        if timestamps is not None and (copy_timestamps or not isinstance(timestamps, list)):
             timestamps = list(timestamps)
 
         self._sample_interval_strategy = sample_interval_strategy
@@ -212,6 +232,24 @@ def __eq__(self, value: object, /) -> bool:
             and self._timestamps == value._timestamps
         )
 
+    def __reduce__(self) -> tuple[Any, ...]:
+        """Return object state for pickling."""
+        ctor_args = (
+            self._sample_interval_mode,
+            self._timestamp,
+            self._time_offset,
+            self._sample_interval,
+            self._timestamps,
+        )
+        ctor_kwargs: dict[str, Any] = {}
+        if self._timestamps is not None:
+            ctor_kwargs["copy_timestamps"] = 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)."""
         # For Enum, __str__ is an unqualified ctor expression like E.V and __repr__ is <E.V: 0>.

src/nitypes/waveform/_timing/_base.py

@@ -14,6 +14,8 @@ class PrecisionTiming(BaseTiming[ht.datetime, ht.timedelta]):
     """High-precision waveform timing using the hightime package.
 
     The hightime package has up to yoctosecond precision.
+
+    Waveform timing objects are immutable.
     """
 
     _DEFAULT_TIME_OFFSET = ht.timedelta()
@@ -66,6 +68,8 @@ def __init__(
         time_offset: ht.timedelta | None = None,
         sample_interval: ht.timedelta | None = None,
         timestamps: Sequence[ht.datetime] | None = None,
+        *,
+        copy_timestamps: bool = True,
     ) -> None:
         """Construct a high-precision waveform timing object.
 
@@ -74,7 +78,14 @@ def __init__(
         - PrecisionTiming.create_with_regular_interval
         - PrecisionTiming.create_with_irregular_interval
         """
-        super().__init__(sample_interval_mode, timestamp, time_offset, sample_interval, timestamps)
+        super().__init__(
+            sample_interval_mode,
+            timestamp,
+            time_offset,
+            sample_interval,
+            timestamps,
+            copy_timestamps=copy_timestamps,
+        )
 
 
 PrecisionTiming.empty = PrecisionTiming.create_with_no_interval()

src/nitypes/waveform/_timing/_precision.py

@@ -14,6 +14,8 @@ class Timing(BaseTiming[dt.datetime, dt.timedelta]):
 
     The standard datetime module has up to microsecond precision. For higher precision, use
     PrecisionTiming.
+
+    Waveform timing objects are immutable.
     """
 
     _DEFAULT_TIME_OFFSET = dt.timedelta()
@@ -66,6 +68,8 @@ def __init__(
         time_offset: dt.timedelta | None = None,
         sample_interval: dt.timedelta | None = None,
         timestamps: Sequence[dt.datetime] | None = None,
+        *,
+        copy_timestamps: bool = True,
     ) -> None:
         """Construct a waveform timing object.
 
@@ -74,7 +78,14 @@ def __init__(
         - Timing.create_with_regular_interval
         - Timing.create_with_irregular_interval
         """
-        super().__init__(sample_interval_mode, timestamp, time_offset, sample_interval, timestamps)
+        super().__init__(
+            sample_interval_mode,
+            timestamp,
+            time_offset,
+            sample_interval,
+            timestamps,
+            copy_timestamps=copy_timestamps,
+        )
 
 
 Timing.empty = Timing.create_with_no_interval()

src/nitypes/waveform/_timing/_standard.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import copy
+import pickle
 from typing import SupportsFloat
 
 import numpy as np
 import numpy.typing as npt
 import pytest
 
 from nitypes._typing import assert_type
-from nitypes.waveform import LinearScaleMode
+from nitypes.waveform import NO_SCALING, LinearScaleMode, ScaleMode
 
 
 @pytest.mark.parametrize(
@@ -80,7 +82,71 @@ def test___float64_ndarray___transform_data___returns_float64_scaled_data() -> N
     assert list(scaled_data) == [4.0, 7.0, 10.0, 13.0]
 
 
+@pytest.mark.parametrize(
+    "left, right",
+    [
+        (LinearScaleMode(1.0, 0.0), LinearScaleMode(1.0, 0.0)),
+        (LinearScaleMode(1.2345, 0.006789), LinearScaleMode(1.2345, 0.006789)),
+    ],
+)
+def test___same_value___equality___equal(left: LinearScaleMode, right: LinearScaleMode) -> None:
+    assert left == right
+    assert not (left != right)
+
+
+@pytest.mark.parametrize(
+    "left, right",
+    [
+        (LinearScaleMode(1.0, 0.0), LinearScaleMode(1.0, 0.1)),
+        (LinearScaleMode(1.0, 0.0), LinearScaleMode(1.1, 0.0)),
+        (LinearScaleMode(1.2345, 0.006789), LinearScaleMode(1.23456, 0.006789)),
+        (LinearScaleMode(1.2345, 0.006789), LinearScaleMode(1.2345, 0.00678)),
+        (LinearScaleMode(1.0, 0.0), NO_SCALING),
+        (NO_SCALING, LinearScaleMode(1.0, 0.0)),
+    ],
+)
+def test___different_value___equality___not_equal(left: ScaleMode, right: ScaleMode) -> None:
+    assert not (left == right)
+    assert left != right
+
+
 def test___scale_mode___repr___looks_ok() -> None:
     scale_mode = LinearScaleMode(1.2345, 0.006789)
 
     assert repr(scale_mode) == "nitypes.waveform.LinearScaleMode(1.2345, 0.006789)"
+
+
+def test___scale_mode___copy___makes_shallow_copy() -> None:
+    scale_mode = LinearScaleMode(1.2345, 0.006789)
+
+    new_scale_mode = copy.copy(scale_mode)
+
+    assert new_scale_mode == scale_mode
+    assert new_scale_mode is not scale_mode
+
+
+def test___scale_mode___deepcopy___makes_deep_copy() -> None:
+    scale_mode = LinearScaleMode(1.2345, 0.006789)
+
+    new_scale_mode = copy.deepcopy(scale_mode)
+
+    assert new_scale_mode == scale_mode
+    assert new_scale_mode is not scale_mode
+
+
+def test___scale_mode___pickle_unpickle___makes_deep_copy() -> None:
+    scale_mode = LinearScaleMode(1.2345, 0.006789)
+
+    new_scale_mode = pickle.loads(pickle.dumps(scale_mode))
+
+    assert new_scale_mode == scale_mode
+    assert new_scale_mode is not scale_mode
+
+
+def test___scale_mode___pickle___references_public_modules() -> None:
+    scale_mode = LinearScaleMode(1.2345, 0.006789)
+
+    scale_mode_bytes = pickle.dumps(scale_mode)
+
+    assert b"nitypes.waveform" in scale_mode_bytes
+    assert b"nitypes.waveform._scaling" not in scale_mode_bytes

tests/unit/waveform/_scaling/test_linear.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import copy
+import pickle
+
 import numpy as np
 import numpy.typing as npt
 
@@ -44,3 +47,31 @@ def test___float64_ndarray___transform_data___returns_float64_scaled_data() -> N
 
 def test___scale_mode___repr___looks_ok() -> None:
     assert repr(NO_SCALING) == "nitypes.waveform.NoneScaleMode()"
+
+
+def test___scale_mode___copy___makes_shallow_copy() -> None:
+    new_scale_mode = copy.copy(NO_SCALING)
+
+    assert new_scale_mode == NO_SCALING
+    assert new_scale_mode is not NO_SCALING
+
+
+def test___scale_mode___deepcopy___makes_deep_copy() -> None:
+    new_scale_mode = copy.deepcopy(NO_SCALING)
+
+    assert new_scale_mode == NO_SCALING
+    assert new_scale_mode is not NO_SCALING
+
+
+def test___scale_mode___pickle_unpickle___makes_deep_copy() -> None:
+    new_scale_mode = pickle.loads(pickle.dumps(NO_SCALING))
+
+    assert new_scale_mode == NO_SCALING
+    assert new_scale_mode is not NO_SCALING
+
+
+def test___scale_mode___pickle___references_public_modules() -> None:
+    scale_mode_bytes = pickle.dumps(NO_SCALING)
+
+    assert b"nitypes.waveform" in scale_mode_bytes
+    assert b"nitypes.waveform._scaling" not in scale_mode_bytes