Handle reads that don't match Waveform length (#819)

* handle short reads

* ReallocationPolicy

* use GROUP_BY_SCAN_NUMBER instead of reshaping read_array

* change ReallocationPolicy to a parameter

* cleanup

* compare against capacity instead of sample_count when using ReallocationPolicy.TO_GROW

* cleanup

* remove dead code

* ReallocationPolicy should default to TO_GROW

* if waveform._start_index + number_of_samples_per_channel > waveform.capacity:

* clean up setting of waveform.sample_count

---------

Co-authored-by: Mike Prosser <Mike.Prosser@emerson.com>

Author: mikeprosserni

generated/nidaqmx/_library_interpreter.py

@@ -6392,6 +6392,8 @@ def read_analog_waveform(
             t0_array = None
             dt_array = None
 
+        waveform.sample_count = number_of_samples_per_channel
+
         error_code, samples_read = self._internal_read_analog_waveform_ex(
             task_handle,
             number_of_samples_per_channel,
@@ -6403,10 +6405,11 @@ def read_analog_waveform(
             dt_array,
         )
 
+        waveform.sample_count = samples_read
+
         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)
         return samples_read
 
@@ -6431,6 +6434,9 @@ def read_analog_waveforms(
             t0_array = None
             dt_array = None
 
+        for waveform in waveforms:
+            waveform.sample_count = number_of_samples_per_channel
+
         error_code, samples_read = self._internal_read_analog_waveform_per_chan(
             task_handle,
             number_of_samples_per_channel,
@@ -6441,10 +6447,12 @@ def read_analog_waveforms(
             dt_array,
         )
 
+        for waveform in waveforms:
+            waveform.sample_count = samples_read
+            
         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)
         return samples_read
 
@@ -6660,6 +6668,8 @@ def read_digital_waveform(
             t0_array = None
             dt_array = None
 
+        waveform.sample_count = number_of_samples_per_channel
+
         error_code, samples_read = self._internal_read_digital_waveform(
             task_handle,
             number_of_samples_per_channel,
@@ -6672,10 +6682,11 @@ def read_digital_waveform(
             None,
         )
 
+        waveform.sample_count = samples_read
+        
         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)
         return samples_read
 
@@ -6705,7 +6716,7 @@ def read_digital_waveforms(
         # Since there's no DAQmxInternalReadDigitalWaveformPerChan, we have to allocate a
         # temporary contiguous array to read the data from multiple channels into.
         read_array = numpy.zeros(
-            (channel_count, number_of_samples_per_channel, number_of_signals_per_sample),
+            (number_of_samples_per_channel, channel_count, number_of_signals_per_sample),
             dtype=numpy.uint8)
 
         bytes_per_chan_array = numpy.zeros(channel_count, dtype=numpy.uint32)
@@ -6714,25 +6725,25 @@ def read_digital_waveforms(
             task_handle,
             number_of_samples_per_channel,
             timeout,
-            FillMode.GROUP_BY_CHANNEL.value,
+            FillMode.GROUP_BY_SCAN_NUMBER.value, # GROUP_BY_SCAN_NUMBER handles short reads better than GROUP_BY_CHANNEL
             read_array,
             properties,
             t0_array,
             dt_array,
             bytes_per_chan_array,
         )
 
-        for i, waveform in enumerate(waveforms):
-            waveform_signals = waveform.data.shape[1]
-            channel_signals = bytes_per_chan_array[i]
-            if waveform_signals != channel_signals:
-                raise ValueError(f"waveforms[{i}].data has {waveform_signals} signals, but expected {channel_signals}")
-            waveform.data[:] = read_array[i, :, :channel_signals]
+        for i, waveform in enumerate(waveforms):        
+            waveform.sample_count = samples_read
+            waveform_signal_count = waveform.data.shape[1]
+            channel_signal_count = bytes_per_chan_array[i]
+            if waveform_signal_count != channel_signal_count:
+                raise ValueError(f"waveforms[{i}].data has {waveform_signal_count} signals, but expected {channel_signal_count}")
+            waveform.data[:] = read_array[:, i, :channel_signal_count]
 
         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)
         return samples_read
 
@@ -6759,7 +6770,7 @@ def read_new_digital_waveforms(
             dt_array = None
 
         read_array = numpy.zeros(
-            (channel_count, number_of_samples_per_channel, number_of_signals_per_sample),
+            (number_of_samples_per_channel, channel_count, number_of_signals_per_sample),
             dtype=numpy.uint8)
 
         bytes_per_chan_array = numpy.zeros(channel_count, dtype=numpy.uint32)
@@ -6768,7 +6779,7 @@ def read_new_digital_waveforms(
             task_handle,
             number_of_samples_per_channel,
             timeout,
-            FillMode.GROUP_BY_CHANNEL.value,
+            FillMode.GROUP_BY_SCAN_NUMBER.value, # GROUP_BY_SCAN_NUMBER handles short reads better than GROUP_BY_CHANNEL
             read_array,
             properties,
             t0_array,
@@ -6778,10 +6789,10 @@ def read_new_digital_waveforms(
 
         waveforms = []
         for i in range(channel_count):
-            signal_count = bytes_per_chan_array[i]
+            channel_signal_count = bytes_per_chan_array[i]
             waveform = DigitalWaveform(
                 sample_count=samples_read,
-                data=read_array[i, :, :signal_count],
+                data=read_array[:, i, :channel_signal_count],
                 copy_extended_properties=False,
                 extended_properties=properties[i] if properties else None)
             waveforms.append(waveform)

generated/nidaqmx/constants.py

@@ -1026,6 +1026,11 @@ class _TriggerUsageTypes(Enum):
     ARM_START = 32  #: Device supports arm start triggers
 
 
+class ReallocationPolicy(Enum):
+    DO_NOT_REALLOCATE = 0  #: Do not reallocate waveforms.
+    TO_GROW = 1  #: Reallocate waveforms to grow when needed.
+
+
 class WaveformAttributeMode(Flag):
     NONE = 0
     TIMING = 1

generated/nidaqmx/stream_readers.py

@@ -4,7 +4,7 @@
 from nidaqmx import DaqError
 
 from nidaqmx._feature_toggles import WAVEFORM_SUPPORT, requires_feature
-from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE
+from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE, ReallocationPolicy
 from nidaqmx.error_codes import DAQmxErrors
 from nitypes.waveform import AnalogWaveform
 
@@ -143,6 +143,7 @@ def read_waveforms(
         self,
         waveforms: list[AnalogWaveform[numpy.float64]],
         number_of_samples_per_channel: int = READ_ALL_AVAILABLE,
+        reallocation_policy: ReallocationPolicy = ReallocationPolicy.TO_GROW,
         timeout: int = 10,
     ) -> int:
         """
@@ -181,6 +182,9 @@ def read_waveforms(
                 "read_all_avail_samp" property to True, the method reads
                 the samples currently available in the buffer and does
                 not wait for the task to acquire all requested samples.
+            reallocation_policy (Optional[ReallocationPolicy]): Specifies
+                the reallocation policy to use when the read yields more
+                samples than the current capacity of the waveform.
             timeout (Optional[float]): Specifies the amount of time in
                 seconds to wait for samples to become available. If the
                 time elapses, the method returns an error and any
@@ -210,13 +214,15 @@ def read_waveforms(
                 DAQmxErrors.MISMATCHED_INPUT_ARRAY_SIZES, task_name=self._task.name)
 
         for i, waveform in enumerate(waveforms):
-            if number_of_samples_per_channel > waveform.sample_count:
-                # TODO: AB#3228924 - if allowed by the caller, increase the sample count of the waveform
-                raise DaqError(
-                    f'The waveform at index {i} does not have enough space ({waveform.sample_count}) to hold '
-                    f'the requested number of samples ({number_of_samples_per_channel}). Please provide larger '
-                    'waveforms or adjust the number of samples requested.',
-                    DAQmxErrors.READ_BUFFER_TOO_SMALL, task_name=self._task.name)
+            if waveform._start_index + number_of_samples_per_channel > waveform.capacity:
+                if reallocation_policy == ReallocationPolicy.TO_GROW:
+                    waveform.capacity = waveform._start_index + number_of_samples_per_channel
+                else:
+                    raise DaqError(
+                        f'The waveform at index {i} does not have enough space ({waveform.capacity - waveform._start_index}) to hold '
+                        f'the requested number of samples ({number_of_samples_per_channel}). Please provide larger '
+                        'waveforms or adjust the number of samples requested.',
+                        DAQmxErrors.READ_BUFFER_TOO_SMALL, task_name=self._task.name)
 
         return self._interpreter.read_analog_waveforms(
             self._handle,

generated/nidaqmx/stream_readers/_analog_multi_channel_reader.py

@@ -4,7 +4,7 @@
 from nidaqmx import DaqError
 
 from nidaqmx._feature_toggles import WAVEFORM_SUPPORT, requires_feature
-from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE
+from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE, ReallocationPolicy
 from nidaqmx.error_codes import DAQmxErrors
 from nitypes.waveform import AnalogWaveform
 
@@ -115,6 +115,7 @@ def read_waveform(
         self, 
         waveform: AnalogWaveform[numpy.float64],
         number_of_samples_per_channel: int = READ_ALL_AVAILABLE, 
+        reallocation_policy: ReallocationPolicy = ReallocationPolicy.TO_GROW,
         timeout: int = 10, 
     ) -> int:
         """
@@ -152,6 +153,9 @@ def read_waveform(
                 "read_all_avail_samp" property to True, the method reads
                 the samples currently available in the buffer and does
                 not wait for the task to acquire all requested samples.
+            reallocation_policy (Optional[ReallocationPolicy]): Specifies
+                the reallocation policy to use when the read yields more
+                samples than the current capacity of the waveform.
             timeout (Optional[float]): Specifies the amount of time in
                 seconds to wait for samples to become available. If the
                 time elapses, the method returns an error and any
@@ -170,12 +174,14 @@ def read_waveform(
             self._task._calculate_num_samps_per_chan(
                 number_of_samples_per_channel))
 
-        if number_of_samples_per_channel > waveform.sample_count:
-            # TODO: AB#3228924 - if allowed by the caller, increase the sample count of the waveform
-            raise DaqError(
-                f'The provided waveform does not have enough space ({waveform.sample_count}) to hold '
-                f'the requested number of samples ({number_of_samples_per_channel}). Please provide a larger '
-                'waveform or adjust the number of samples requested.',
+        if waveform._start_index + number_of_samples_per_channel > waveform.capacity:
+            if reallocation_policy == ReallocationPolicy.TO_GROW:
+                waveform.capacity = waveform._start_index + number_of_samples_per_channel
+            else:
+                raise DaqError(
+                    f'The provided waveform does not have enough space ({waveform.capacity - waveform._start_index}) to hold '
+                    f'the requested number of samples ({number_of_samples_per_channel}). Please provide a larger '
+                    'waveform or adjust the number of samples requested.',
                 DAQmxErrors.READ_BUFFER_TOO_SMALL, task_name=self._task.name)
 
         return self._interpreter.read_analog_waveform(

generated/nidaqmx/stream_readers/_analog_single_channel_reader.py

@@ -4,7 +4,7 @@
 from nidaqmx import DaqError
 
 from nidaqmx._feature_toggles import WAVEFORM_SUPPORT, requires_feature
-from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE
+from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE, ReallocationPolicy
 from nidaqmx.error_codes import DAQmxErrors
 from nitypes.waveform import DigitalWaveform
 
@@ -487,6 +487,7 @@ def read_waveforms(
         self,
         waveforms: list[DigitalWaveform[numpy.uint8]],
         number_of_samples_per_channel: int = READ_ALL_AVAILABLE,
+        reallocation_policy: ReallocationPolicy = ReallocationPolicy.TO_GROW,
         timeout: int = 10,
     ) -> list[DigitalWaveform[numpy.uint8]]:
         """
@@ -518,6 +519,9 @@ def read_waveforms(
                 "read_all_avail_samp" property to True, the method reads
                 the samples currently available in the buffer and does
                 not wait for the task to acquire all requested samples.
+            reallocation_policy (Optional[ReallocationPolicy]): Specifies
+                the reallocation policy to use when the read yields more
+                samples than the current capacity of the waveform.
             timeout (Optional[float]): Specifies the amount of time in
                 seconds to wait for samples to become available. If the
                 time elapses, the method returns an error and any
@@ -547,13 +551,15 @@ def read_waveforms(
                 DAQmxErrors.MISMATCHED_INPUT_ARRAY_SIZES, task_name=self._task.name)
 
         for i, waveform in enumerate(waveforms):
-            if number_of_samples_per_channel > waveform.sample_count:
-                # TODO: AB#3228924 - if allowed by the caller, increase the sample count of the waveform
-                raise DaqError(
-                    f'The waveform at index {i} does not have enough space ({waveform.sample_count}) to hold '
-                    f'the requested number of samples ({number_of_samples_per_channel}). Please provide larger '
-                    'waveforms or adjust the number of samples requested.',
-                    DAQmxErrors.READ_BUFFER_TOO_SMALL, task_name=self._task.name)
+            if waveform._start_index + number_of_samples_per_channel > waveform.capacity:
+                if reallocation_policy == ReallocationPolicy.TO_GROW:
+                    waveform.capacity = waveform._start_index + number_of_samples_per_channel
+                else:
+                    raise DaqError(
+                        f'The waveform at index {i} does not have enough space ({waveform.capacity - waveform._start_index}) to hold '
+                        f'the requested number of samples ({number_of_samples_per_channel}). Please provide larger '
+                        'waveforms or adjust the number of samples requested.',
+                        DAQmxErrors.READ_BUFFER_TOO_SMALL, task_name=self._task.name)
 
         waveforms = self._interpreter.read_digital_waveforms(
             self._handle,

generated/nidaqmx/stream_readers/_digital_multi_channel_reader.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import numpy
+from nidaqmx import DaqError
 
 from nidaqmx._feature_toggles import WAVEFORM_SUPPORT, requires_feature
-from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE
+from nidaqmx.constants import FillMode, READ_ALL_AVAILABLE, ReallocationPolicy
+from nidaqmx.error_codes import DAQmxErrors
 from nitypes.waveform import DigitalWaveform
 
 from nidaqmx.stream_readers._channel_reader_base import ChannelReaderBase
@@ -388,6 +390,7 @@ def read_waveform(
         self,
         waveform: DigitalWaveform[numpy.uint8],
         number_of_samples_per_channel: int = READ_ALL_AVAILABLE,
+        reallocation_policy: ReallocationPolicy = ReallocationPolicy.TO_GROW,
         timeout: float = 10.0,
     ) -> int:
         """
@@ -418,6 +421,9 @@ def read_waveform(
                 "read_all_avail_samp" property to True, the method reads
                 the samples currently available in the buffer and does
                 not wait for the task to acquire all requested samples.
+            reallocation_policy (Optional[ReallocationPolicy]): Specifies
+                the reallocation policy to use when the read yields more
+                samples than the current capacity of the waveform.
             timeout (Optional[float]): Specifies the amount of time in
                 seconds to wait for samples to become available. If the
                 time elapses, the method returns an error and any
@@ -438,6 +444,16 @@ def read_waveform(
             self._task._calculate_num_samps_per_chan(
                 number_of_samples_per_channel))
 
+        if waveform._start_index + number_of_samples_per_channel > waveform.capacity:
+            if reallocation_policy == ReallocationPolicy.TO_GROW:
+                waveform.capacity = waveform._start_index + number_of_samples_per_channel
+            else:
+                raise DaqError(
+                    f'The waveform does not have enough space ({waveform.capacity - waveform._start_index}) to hold '
+                    f'the requested number of samples ({number_of_samples_per_channel}). Please '
+                    'provide a larger waveform or adjust the number of samples requested.',
+                    DAQmxErrors.READ_BUFFER_TOO_SMALL, task_name=self._task.name)
+    
         return self._interpreter.read_digital_waveform(
             self._handle, 
             number_of_samples_per_channel,