Support writing directly with wrsamp.

Author: tompollard

tests/test_archive.py

@@ -98,7 +98,7 @@ def test_wfdb_archive_inline_round_trip():
         sig = (np.random.randn(sig_len, 2) * 1000).astype(np.float32)
 
         # Create archive inline using context manager
-        with WFDBArchive(record_basename, mode="w") as wfdb_archive:
+        with WFDBArchive(record_path, mode="w") as wfdb_archive:
             wrsamp(
                 record_name=record_basename,
                 fs=fs,

wfdb/io/_signal.py

@@ -1710,6 +1710,16 @@ def _rd_dat_file(
 
     # Local file or .wfdb archive
     if wfdb_archive is not None:
+        # If the exact file name isn't found, look for any .dat file
+        if not wfdb_archive.exists(file_name):
+            dat_files = [f for f in wfdb_archive.zipfile.namelist()
+                        if f.endswith('.dat')]
+            if not dat_files:
+                raise FileNotFoundError(
+                    f"No dat file found in archive for {file_name}"
+                )
+            file_name = dat_files[0]  # Use the first dat file found
+
         with wfdb_archive.open(file_name, "rb") as fp:
             fp.seek(start_byte)
             sig_data = util.fromfile(

wfdb/io/archive.py

@@ -31,7 +31,11 @@ def __init__(self, record_name, mode="r"):
             'r' for read (default), 'w' for write.
         """
         self.record_name = record_name
-        self.archive_path = f"{record_name}.wfdb"
+        # Only append .wfdb if it's not already there
+        if record_name.endswith('.wfdb'):
+            self.archive_path = record_name
+        else:
+            self.archive_path = f"{record_name}.wfdb"
         self.zipfile = None
         self.mode = mode
 
@@ -41,14 +45,24 @@ def __init__(self, record_name, mode="r"):
                     f"Archive not found: {self.archive_path}"
                 )
             if not zipfile.is_zipfile(self.archive_path):
-                raise ValueError(f"Invalid WFDB archive: {self.archive_path}")
-            self.zipfile = zipfile.ZipFile(self.archive_path, mode="r")
+                raise ValueError(
+                    f"Invalid WFDB archive: {self.archive_path}"
+                )
+            self.zipfile = zipfile.ZipFile(
+                self.archive_path, mode="r"
+            )
 
         elif mode == "w":
             # Create archive file if needed
             if not os.path.exists(self.archive_path):
-                WFDBArchive.make_archive_file([], self.archive_path)
-            self.zipfile = zipfile.ZipFile(self.archive_path, mode="a")
+                # Create the directory if it doesn't exist
+                os.makedirs(os.path.dirname(os.path.abspath(self.archive_path)), exist_ok=True)
+                WFDBArchive.make_archive_file(
+                    [], self.archive_path
+                )
+            self.zipfile = zipfile.ZipFile(
+                self.archive_path, mode="a"
+            )
 
     def __enter__(self):
         return self
@@ -85,7 +99,8 @@ def open(self, filename, mode="r"):
                     yield io.TextIOWrapper(f)
         else:
             raise FileNotFoundError(
-                f"Could not find '{filename}' as loose file or inside '{self.archive_path}'."
+                f"Could not find '{filename}' as loose file or inside "
+                f"'{self.archive_path}'."
             )
 
     def close(self):

wfdb/io/record.py

@@ -10,7 +10,7 @@
 
 from wfdb.io import _header, _signal, _url, download, header, util
 from wfdb.io._coreio import CLOUD_PROTOCOLS
-from wfdb.io.archive import get_archive
+from wfdb.io.archive import get_archive, WFDBArchive
 
 # -------------- WFDB Signal Calibration and Classification ---------- #
 
@@ -911,6 +911,10 @@ def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None):
             of the uniform signal (d_signal).
         write_dir : str, optional
             The directory in which to write the files.
+        wfdb_archive : str or WFDBArchive, optional
+            If provided, writes the record to a .wfdb archive. Can be either:
+            - A string path to the archive file (e.g., 'record.wfdb')
+            - A WFDBArchive object for more advanced usage
 
         Returns
         -------
@@ -927,6 +931,15 @@ def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None):
                     checksums[ch] = old_val
             self.checksum = checksums
 
+        # Handle wfdb_archive parameter
+        if wfdb_archive:
+            if isinstance(wfdb_archive, str):
+                # If a string path is provided, create a WFDBArchive object
+                from wfdb.io.archive import get_archive
+                wfdb_archive = get_archive(wfdb_archive, mode="w")
+            elif not isinstance(wfdb_archive, WFDBArchive):
+                raise TypeError("wfdb_archive must be either a string path or WFDBArchive object")
+
         # Perform field validity and cohesion checks, and write the
         # header file.
         self.wrheader(
@@ -1975,7 +1988,6 @@ def rdrecord(
         Option used to stream data from Physionet. The Physionet
         database directory from which to find the required record files.
         eg. For record '100' in 'http://physionet.org/content/mitdb'
-        pn_dir='mitdb'.
     m2s : bool, optional
         Used when reading multi-segment records. Specifies whether to
         directly return a WFDB MultiRecord object (False), or to convert
@@ -2033,15 +2045,19 @@ def rdrecord(
     --------
     >>> record = wfdb.rdrecord('sample-data/test01_00s', sampfrom=800,
                                channels=[1, 3])
-
     """
     wfdb_archive = None
     is_wfdb_archive = record_name.endswith(".wfdb")
 
     if is_wfdb_archive:
         record_base = record_name[:-5]  # remove ".wfdb"
         wfdb_archive = get_archive(record_base)
-        hea_file = os.path.basename(record_base) + ".hea"
+
+        # Find any .hea file in the archive
+        hea_files = [f for f in wfdb_archive.zipfile.namelist() if f.endswith('.hea')]
+        if not hea_files:
+            raise FileNotFoundError(f"No header file found in archive {record_name}")
+        hea_file = hea_files[0]  # Use the first header file found
 
         import tempfile
 
@@ -2954,6 +2970,10 @@ def wrsamp(
         setting both `base_date` and `base_time`.
     write_dir : str, optional
         The directory in which to write the files.
+    wfdb_archive : str or WFDBArchive, optional
+        If provided, writes the record to a .wfdb archive. Can be either:
+        - A string path to the archive file (e.g., 'record.wfdb')
+        - A WFDBArchive object for more advanced usage
 
     Returns
     -------
@@ -2978,6 +2998,10 @@ def wrsamp(
     >>> # Write a local WFDB record (manually inserting fields)
     >>> wfdb.wrsamp('ecgrecord', fs = 250, units=['mV', 'mV'],
                     sig_name=['I', 'II'], p_signal=signals, fmt=['16', '16'])
+    >>> # Write to a .wfdb archive using a string path
+    >>> wfdb.wrsamp('ecgrecord', fs = 250, units=['mV', 'mV'],
+                    sig_name=['I', 'II'], p_signal=signals, fmt=['16', '16'],
+                    wfdb_archive='ecgrecord.wfdb')
 
     """
     # Check for valid record name
@@ -3082,10 +3106,14 @@ def wrsamp(
     else:
         expanded = False
 
+    # Handle wfdb_archive parameter
     if wfdb_archive:
-        wfdb_archive = get_archive(
-            os.path.join(write_dir, record_name), mode="w"
-        )
+        if isinstance(wfdb_archive, str):
+            # If a string path is provided, create a WFDBArchive object
+            from wfdb.io.archive import get_archive
+            wfdb_archive = get_archive(wfdb_archive, mode="w")
+        elif not isinstance(wfdb_archive, WFDBArchive):
+            raise TypeError("wfdb_archive must be either a string path or WFDBArchive object")
 
     # Write the record files - header and associated dat
     record.wrsamp(