Add in_stream.waveform_attribute_mode (#800)

* first draft

* cleanup

* test cleanup

* _internal_read_analog_waveform_ex to take in Nones to avoid allocations

* fix merge

* implement WaveformAttributeMode in read_analog_waveforms and internal_read_analog_waveform_per_chan

* refactor how the arrays for timing information are handled

* _get_wfm_attr_callback() helper

* cleanup

---------

Co-authored-by: Mike Prosser <[email protected]>

Author: mikeprosserni

generated/nidaqmx/_base_interpreter.py

@@ -3,6 +3,7 @@
 import numpy
 from nitypes.waveform import AnalogWaveform
 from typing import Sequence
+from nidaqmx.constants import WaveformAttributeMode
 
 
 class BaseEventHandler(abc.ABC):
@@ -1855,7 +1856,8 @@ def read_analog_waveform(
         task_handle: object,
         number_of_samples_per_channel: int,
         timeout: float,
-        waveform: AnalogWaveform[numpy.float64]
+        waveform: AnalogWaveform[numpy.float64],
+        waveform_attribute_mode: WaveformAttributeMode
     ) -> None:
         raise NotImplementedError
 
@@ -1865,6 +1867,7 @@ def read_analog_waveforms(
         task_handle: object,
         number_of_samples_per_channel: int,
         timeout: float,
-        waveforms: Sequence[AnalogWaveform[numpy.float64]]
+        waveforms: Sequence[AnalogWaveform[numpy.float64]],
+        waveform_attribute_mode: WaveformAttributeMode
     ) -> None:
         raise NotImplementedError

generated/nidaqmx/_grpc_interpreter.py

@@ -18,6 +18,7 @@
 from nidaqmx._stubs import nidaqmx_pb2 as grpc_types
 from nidaqmx._stubs import nidaqmx_pb2_grpc as nidaqmx_grpc
 from nidaqmx._stubs import session_pb2 as session_grpc_types
+from nidaqmx.constants import WaveformAttributeMode
 from nidaqmx.error_codes import DAQmxErrors
 from nidaqmx._grpc_time import convert_time_to_timestamp, convert_timestamp_to_time
 
@@ -3606,7 +3607,8 @@ def read_analog_waveform(
         task_handle: object,
         number_of_samples_per_channel: int,
         timeout: float,
-        waveform: AnalogWaveform[numpy.float64]
+        waveform: AnalogWaveform[numpy.float64],
+        waveform_attribute_mode: WaveformAttributeMode
     ) -> None:
         raise NotImplementedError
 
@@ -3615,7 +3617,8 @@ def read_analog_waveforms(
         task_handle: object,
         number_of_samples_per_channel: int,
         timeout: float,
-        waveforms: Sequence[AnalogWaveform[numpy.float64]]
+        waveforms: Sequence[AnalogWaveform[numpy.float64]],
+        waveform_attribute_mode: WaveformAttributeMode
     ) -> None:
         raise NotImplementedError
 

generated/nidaqmx/_library_interpreter.py

@@ -11,11 +11,11 @@
 from datetime import timezone
 from hightime import datetime as ht_datetime
 from hightime import timedelta as ht_timedelta
-from typing import Callable, List, Optional, Sequence, Tuple, TYPE_CHECKING
+from typing import Callable, List, MutableSequence, Optional, Sequence, Tuple, TYPE_CHECKING
 
 from nidaqmx._base_interpreter import BaseEventHandler, BaseInterpreter
 from nidaqmx._lib import lib_importer, ctypes_byte_str, c_bool32, wrapped_ndpointer, TaskHandle
-from nidaqmx.constants import FillMode
+from nidaqmx.constants import FillMode, WaveformAttributeMode
 from nidaqmx.error_codes import DAQmxErrors, DAQmxWarnings
 from nidaqmx.errors import DaqError, DaqFunctionNotSupportedError, DaqReadError, DaqWarning, DaqWriteError
 from nidaqmx._lib_time import AbsoluteTime
@@ -6376,24 +6376,35 @@ def read_analog_waveform(
         task_handle: object,
         number_of_samples_per_channel: int,
         timeout: float,
-        waveform: AnalogWaveform[numpy.float64]
+        waveform: AnalogWaveform[numpy.float64],
+        waveform_attribute_mode: WaveformAttributeMode
     ) -> None:
         """Read an analog waveform with timing and attributes."""
-        error_code, samples_read, timestamps, sample_intervals = self._internal_read_analog_waveform_ex(
+        if WaveformAttributeMode.EXTENDED_PROPERTIES in waveform_attribute_mode:
+            properties = [waveform.extended_properties]
+        else:
+            properties = None
+
+        if WaveformAttributeMode.TIMING in waveform_attribute_mode:
+            t0_array = numpy.zeros(1, dtype=numpy.int64)
+            dt_array = numpy.zeros(1, dtype=numpy.int64)
+        else:
+            t0_array = None
+            dt_array = None
+
+        error_code, samples_read = self._internal_read_analog_waveform_ex(
             task_handle,
-            1, # single channel
             number_of_samples_per_channel,
             timeout,
             FillMode.GROUP_BY_CHANNEL.value,
             waveform.raw_data,
-            [waveform.extended_properties]
+            properties,
+            t0_array,
+            dt_array,
         )
 
-        waveform.timing = Timing(
-            sample_interval_mode=SampleIntervalMode.REGULAR,
-            timestamp=timestamps[0],
-            sample_interval=sample_intervals[0],
-        )
+        if t0_array is not None and dt_array is not None:
+            self._set_waveform_timings([waveform], t0_array, dt_array)
 
         # TODO: AB#3228924 - if the read was short, set waveform.sample_count before throwing the exception
         self.check_for_error(error_code, samps_per_chan_read=samples_read)
@@ -6403,41 +6414,51 @@ def read_analog_waveforms(
         task_handle: object,
         number_of_samples_per_channel: int,
         timeout: float,
-        waveforms: Sequence[AnalogWaveform[numpy.float64]]
+        waveforms: Sequence[AnalogWaveform[numpy.float64]],
+        waveform_attribute_mode: WaveformAttributeMode
     ) -> None:
         """Read a set of analog waveforms with timing and attributes. All of the waveforms must be the same size."""
-        error_code, samples_read, timestamps, sample_intervals = self.internal_read_analog_waveform_per_chan(
+        if WaveformAttributeMode.EXTENDED_PROPERTIES in waveform_attribute_mode:
+            properties = [waveform.extended_properties for waveform in waveforms]
+        else:
+            properties = None
+
+        if WaveformAttributeMode.TIMING in waveform_attribute_mode:
+            t0_array = numpy.zeros(len(waveforms), dtype=numpy.int64)
+            dt_array = numpy.zeros(len(waveforms), dtype=numpy.int64)
+        else:
+            t0_array = None
+            dt_array = None
+
+        error_code, samples_read = self.internal_read_analog_waveform_per_chan(
             task_handle,
             number_of_samples_per_channel,
             timeout,
             [waveform.raw_data for waveform in waveforms],
-            [waveform.extended_properties for waveform in waveforms]
+            properties,
+            t0_array,
+            dt_array,
         )
 
-        for i, waveform in enumerate(waveforms):
-            waveform.timing = Timing(
-                sample_interval_mode=SampleIntervalMode.REGULAR,
-                timestamp=timestamps[i],
-                sample_interval=sample_intervals[i],
-            )
+        if t0_array is not None and dt_array is not None:
+            self._set_waveform_timings(waveforms, t0_array, dt_array)
 
         # TODO: AB#3228924 - if the read was short, set waveform.sample_count before throwing the exception
         self.check_for_error(error_code, samps_per_chan_read=samples_read)
 
     def _internal_read_analog_waveform_ex(
         self,
         task_handle: object,
-        channel_count: int,
         number_of_samples_per_channel: int,
         timeout: float,
         fill_mode: int,
         read_array: numpy.typing.NDArray[numpy.float64],
-        properties: Sequence[ExtendedPropertyDictionary]
+        properties: Sequence[ExtendedPropertyDictionary] | None,
+        t0_array: numpy.typing.NDArray[numpy.int64] | None,
+        dt_array: numpy.typing.NDArray[numpy.int64] | None,
     ) -> Tuple[
         int, # error code
         int, # The number of samples per channel that were read
-        Sequence[ht_datetime], # The timestamps for each sample, indexed by channel
-        Sequence[ht_timedelta], # The sample intervals, indexed by channel
     ]:
         assert isinstance(task_handle, TaskHandle)
         samps_per_chan_read = ctypes.c_int()
@@ -6462,19 +6483,6 @@ def _internal_read_analog_waveform_ex(
                         ctypes.POINTER(c_bool32),
                     ]
 
-        t0_array = numpy.zeros(channel_count, dtype=numpy.int64)
-        dt_array = numpy.zeros(channel_count, dtype=numpy.int64)
-
-        def set_wfm_attr_callback(
-            channel_index: int,
-            attribute_name: str,
-            attribute_type: WfmAttrType,
-            value: ExtendedPropertyValue,
-            callback_data: object,
-        ) -> int:
-            properties[channel_index][attribute_name] = value
-            return 0
-
         error_code = cfunc(
             task_handle,
             number_of_samples_per_channel,
@@ -6483,31 +6491,28 @@ def set_wfm_attr_callback(
             t0_array,
             dt_array,
             0 if t0_array is None else t0_array.size,
-            self._get_wfm_attr_callback_ptr(set_wfm_attr_callback),
+            self._get_wfm_attr_callback(properties),
             None,
             read_array,
             read_array.size,
             ctypes.byref(samps_per_chan_read),
             None,
         )
 
-        timestamps = [_T0_EPOCH + ht_timedelta(seconds=t0 * _INT64_WFM_SEC_PER_TICK) for t0 in t0_array]
-        sample_intervals = [ht_timedelta(seconds=dt * _INT64_WFM_SEC_PER_TICK) for dt in dt_array]
-
-        return error_code, samps_per_chan_read.value, timestamps, sample_intervals
+        return error_code, samps_per_chan_read.value
 
     def internal_read_analog_waveform_per_chan(
         self,
         task_handle: object,
         num_samps_per_chan: int,
         timeout: float,
         read_arrays: Sequence[numpy.typing.NDArray[numpy.float64]],
-        properties: Sequence[ExtendedPropertyDictionary]
+        properties: Sequence[ExtendedPropertyDictionary] | None,
+        t0_array: numpy.typing.NDArray[numpy.int64] | None,
+        dt_array: numpy.typing.NDArray[numpy.int64] | None,
     ) -> Tuple[
         int, # error code
         int, # The number of samples per channel that were read
-        Sequence[ht_datetime], # The timestamps for each sample, indexed by channel
-        Sequence[ht_timedelta], # The sample intervals, indexed by channel
     ]:
         assert isinstance(task_handle, TaskHandle)
         samps_per_chan_read = ctypes.c_int()
@@ -6537,31 +6542,18 @@ def internal_read_analog_waveform_per_chan(
                         ctypes.POINTER(c_bool32),
                     ]
 
-        t0_array = numpy.zeros(channel_count, dtype=numpy.int64)
-        dt_array = numpy.zeros(channel_count, dtype=numpy.int64)
-
         read_array_ptrs = (ctypes.POINTER(ctypes.c_double) * channel_count)()
         for i, read_array in enumerate(read_arrays):
             read_array_ptrs[i] = read_array.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
 
-        def set_wfm_attr_callback(
-            channel_index: int,
-            attribute_name: str,
-            attribute_type: WfmAttrType,
-            value: ExtendedPropertyValue,
-            callback_data: object,
-        ) -> int:
-            properties[channel_index][attribute_name] = value
-            return 0
-
         error_code = cfunc(
             task_handle,
             num_samps_per_chan,
             timeout,
             t0_array,
             dt_array,
             0 if t0_array is None else t0_array.size,
-            self._get_wfm_attr_callback_ptr(set_wfm_attr_callback),
+            self._get_wfm_attr_callback(properties),
             None,
             read_array_ptrs,
             channel_count,
@@ -6571,11 +6563,22 @@ def set_wfm_attr_callback(
         )
         self.check_for_error(error_code, samps_per_chan_read=samps_per_chan_read.value)
 
-        timestamps = [_T0_EPOCH + ht_timedelta(seconds=t0 * _INT64_WFM_SEC_PER_TICK) for t0 in t0_array]
-        sample_intervals = [ht_timedelta(seconds=dt * _INT64_WFM_SEC_PER_TICK) for dt in dt_array]
-
-        return error_code, samps_per_chan_read.value, timestamps, sample_intervals
-
+        return error_code, samps_per_chan_read.value
+
+    def _get_wfm_attr_callback(self, properties):
+        if properties is not None:
+            def set_wfm_attr_callback(
+                channel_index: int,
+                attribute_name: str,
+                attribute_type: WfmAttrType,
+                value: ExtendedPropertyValue,
+                callback_data: object,
+            ) -> int:
+                properties[channel_index][attribute_name] = value
+                return 0
+            return self._get_wfm_attr_callback_ptr(set_wfm_attr_callback)
+        else:
+            return CSetWfmAttrCallbackPtr()
 
     def _get_wfm_attr_value(
         self, attribute_type: int, value: ctypes.c_void_p, value_size_in_bytes: int
@@ -6597,11 +6600,8 @@ def _get_wfm_attr_value(
             raise ValueError(f"Unsupported attribute type {attribute_type}")
 
     def _get_wfm_attr_callback_ptr(
-        self, set_wfm_attr_callback: Optional[SetWfmAttrCallback]
+        self, set_wfm_attr_callback: SetWfmAttrCallback
     ) -> ctypes._FuncPointer:
-        if set_wfm_attr_callback is None:
-            return CSetWfmAttrCallbackPtr()
-
         def _invoke_callback(
             channel_index: int,
             attribute_name: bytes,
@@ -6624,6 +6624,19 @@ def _invoke_callback(
 
         return CSetWfmAttrCallbackPtr(_invoke_callback)
 
+    def _set_waveform_timings(
+        self, 
+        waveforms: Sequence[AnalogWaveform[numpy.float64]], 
+        t0_array: numpy.typing.NDArray[numpy.int64], 
+        dt_array: numpy.typing.NDArray[numpy.int64]
+    ) -> None:
+        for i, waveform in enumerate(waveforms):
+            waveform.timing = Timing(
+                sample_interval_mode=SampleIntervalMode.REGULAR,
+                timestamp=_T0_EPOCH + ht_timedelta(seconds=t0_array[i] * _INT64_WFM_SEC_PER_TICK),
+                sample_interval=ht_timedelta(seconds=dt_array[i] * _INT64_WFM_SEC_PER_TICK),
+            )
+
     def read_id_pin_memory(self, device_name, id_pin_name):
         data_length_read = ctypes.c_uint()
         format_code = ctypes.c_uint()

generated/nidaqmx/constants.py

@@ -1,6 +1,6 @@
 # Do not edit this file; it was automatically generated.
 
-from enum import Enum
+from enum import Enum, Flag
 
 # Constants
 AUTO = -1
@@ -1026,3 +1026,8 @@ class _TriggerUsageTypes(Enum):
     ARM_START = 32  #: Device supports arm start triggers
 
 
+class WaveformAttributeMode(Flag):
+    NONE = 0
+    TIMING = 1
+    EXTENDED_PROPERTIES = 2
+

generated/nidaqmx/stream_readers.py

@@ -317,6 +317,7 @@ def read_waveform(
             number_of_samples_per_channel,
             timeout,
             waveform,
+            self._in_stream.waveform_attribute_mode,
         )
 
         return waveform
@@ -543,6 +544,7 @@ def read_waveforms(
             number_of_samples_per_channel,
             timeout,
             waveforms,
+            self._in_stream.waveform_attribute_mode,
         )
 
         return waveforms