Add validation to MultiRecord.wrsamp for segment cohesion

changelog.d/multiseg-write-validation.rst

@@ -0,0 +1,2 @@
+**Validate multi-segment record writes**
+  ``MultiRecord.wrsamp`` now checks the segment lengths and sampling frequencies of all segments before writing any files, instead of writing a segment whose length does not match ``seg_len``.

src/wfdb/io/record.py

@@ -1192,6 +1192,9 @@ def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None):
         N/A
 
         """
+        # Validate segment cohesion before writing anything.
+        self._check_segment_cohesion()
+
         with archive_for_writing(wfdb_archive) as archive:
             # Perform field validity and cohesion checks, and write the
             # header file.
@@ -1221,7 +1224,14 @@ def _check_segment_cohesion(self):
         if self.n_seg != len(self.segments):
             raise ValueError("Length of segments must match the 'n_seg' field")
 
+        if len(self.seg_len) != self.n_seg:
+            raise ValueError("Length of seg_len must match the 'n_seg' field")
+
         for seg_num, segment in enumerate(self.segments):
+            # Empty segments ('~') have no Record to check
+            if segment is None:
+                continue
+
             # If segment 0 is a layout specification record, check that its file names
             # are all == '~''
             if seg_num == 0 and self.seg_len[0] == 0:
@@ -1241,7 +1251,8 @@ def _check_segment_cohesion(self):
             # field
             if segment.sig_len != self.seg_len[seg_num]:
                 raise ValueError(
-                    f"The signal length of segment {seg_num} does not match the corresponding segment length"
+                    f"The signal length of segment {seg_num} does not match the "
+                    f"corresponding segment length"
                 )
 
         # No need to check the sum of sig_lens from each segment object against sig_len

tests/test_multi_record.py

@@ -1,3 +1,6 @@
+import numpy as np
+import pytest
+
 import wfdb
 
 
@@ -105,3 +108,73 @@ def test_contained_combined_ranges_all_pairs(self):
                     continue
 
                 assert record.contained_combined_ranges([sig_1, sig_2]) == [(0, 2000)]
+
+
+class TestMultiRecordWriting:
+    """Test writing multi-segment records."""
+
+    record_path = "sample-data/multi-segment/041s/041s"
+
+    def test_round_trip(self, tmp_path):
+        record = wfdb.rdrecord(self.record_path, m2s=False, physical=False)
+        expected_signals = [segment.d_signal.copy() for segment in record.segments]
+
+        record.wrsamp(write_dir=str(tmp_path))
+
+        record_write = wfdb.rdrecord(str(tmp_path / "041s"), m2s=False, physical=False)
+        assert record_write.seg_len == record.seg_len
+        assert record_write.sig_name == record.sig_name
+        for segment, expected_signal in zip(record_write.segments, expected_signals):
+            np.testing.assert_array_equal(segment.d_signal, expected_signal)
+
+    def test_segment_length_mismatch_writes_no_files(self, tmp_path):
+        record = wfdb.rdrecord(self.record_path, m2s=False, physical=False)
+        record.seg_len[1] -= 1
+
+        with pytest.raises(ValueError, match="signal length of segment 1"):
+            record.wrsamp(write_dir=str(tmp_path))
+
+        assert list(tmp_path.iterdir()) == []
+
+    def test_short_seg_len_writes_no_files(self, tmp_path):
+        record = wfdb.rdrecord(self.record_path, m2s=False, physical=False)
+        record.seg_len.pop()
+
+        with pytest.raises(
+            ValueError, match="Length of seg_len must match the 'n_seg' field"
+        ):
+            record.wrsamp(write_dir=str(tmp_path))
+
+        assert list(tmp_path.iterdir()) == []
+
+    def test_outdated_expanded_signal_is_ignored_when_not_expanded(self, tmp_path):
+        record = wfdb.rdrecord(self.record_path, m2s=False, physical=False)
+        expected_signal = record.segments[0].d_signal.copy()
+        record.segments[0].e_d_signal = [[] for _ in range(record.segments[0].n_sig)]
+
+        record.wrsamp(write_dir=str(tmp_path), expanded=False)
+
+        record_write = wfdb.rdrecord(str(tmp_path / "041s"), m2s=False, physical=False)
+        np.testing.assert_array_equal(
+            record_write.segments[0].d_signal, expected_signal
+        )
+
+    def test_unspecified_samps_per_frame_is_accepted_when_expanded(self, tmp_path):
+        record = wfdb.rdrecord(
+            self.record_path,
+            m2s=False,
+            physical=False,
+            smooth_frames=False,
+            channels=[6],
+        )
+        record.segments[0].samps_per_frame[0] = None
+
+        record.wrsamp(write_dir=str(tmp_path), expanded=True)
+
+        record_write = wfdb.rdrecord(
+            str(tmp_path / "041s"),
+            m2s=False,
+            physical=False,
+            smooth_frames=False,
+        )
+        assert record_write.segments[0].samps_per_frame == [1]