write file

Author: tompollard

wfdb/io/_signal.py

@@ -1,4 +1,5 @@
 import math
+import io
 import os
 import posixpath
 import sys
@@ -120,7 +121,7 @@ class SignalMixin(object):
 
     """
 
-    def wr_dats(self, expanded, write_dir):
+    def wr_dats(self, expanded, write_dir, wfdb_archive=None):
         """
         Write all dat files associated with a record
         expanded=True to use e_d_signal instead of d_signal.
@@ -132,6 +133,8 @@ def wr_dats(self, expanded, write_dir):
             the `d_signal` attribute (False).
         write_dir : str
             The directory to write the output file to.
+        wfdb_archive : WFDBArchive, optional
+            If set, writes to a .wfdb archive instead of the filesystem.
 
         Returns
         -------
@@ -160,7 +163,8 @@ def wr_dats(self, expanded, write_dir):
         self.check_sig_cohesion([], expanded)
 
         # Write each of the specified dat files
-        self.wr_dat_files(expanded=expanded, write_dir=write_dir)
+        self.wr_dat_files(expanded=expanded, write_dir=write_dir,
+                          wfdb_archive=wfdb_archive)
 
     def check_sig_cohesion(self, write_fields, expanded):
         """
@@ -958,7 +962,7 @@ def calc_checksum(self, expanded=False):
             cs = [int(c) for c in cs]
         return cs
 
-    def wr_dat_files(self, expanded=False, write_dir=""):
+    def wr_dat_files(self, expanded=False, write_dir="", wfdb_archive=None):
         """
         Write each of the specified dat files.
 
@@ -969,6 +973,8 @@ def wr_dat_files(self, expanded=False, write_dir=""):
             the `d_signal` attribute (False).
         write_dir : str, optional
             The directory to write the output file to.
+        wfdb_archive : WFDBArchive, optional
+            If set, writes to a .wfdb archive instead of the local filesystem.
 
         Returns
         -------
@@ -1003,6 +1009,7 @@ def wr_dat_files(self, expanded=False, write_dir=""):
                     [self.e_d_signal[ch] for ch in dat_channels[fn]],
                     [self.samps_per_frame[ch] for ch in dat_channels[fn]],
                     write_dir=write_dir,
+                    wfdb_archive=wfdb_archive,
                 )
         else:
             dsig = self.d_signal
@@ -1013,6 +1020,7 @@ def wr_dat_files(self, expanded=False, write_dir=""):
                     dsig[:, dat_channels[fn][0] : dat_channels[fn][-1] + 1],
                     dat_offsets[fn],
                     write_dir=write_dir,
+                    wfdb_archive=wfdb_archive,
                 )
 
     def smooth_frames(self, sigtype="physical"):
@@ -2322,6 +2330,7 @@ def wr_dat_file(
     e_d_signal=None,
     samps_per_frame=None,
     write_dir="",
+    wfdb_archive=None,
 ):
     """
     Write a dat file. All bytes are written one at a time to avoid
@@ -2519,16 +2528,30 @@ def wr_dat_file(
         else:
             raise ValueError(f"unknown format ({fmt})")
 
-        sf = soundfile.SoundFile(
-            file_path,
-            mode="w",
-            samplerate=96000,
-            channels=n_sig,
-            subtype=subtype,
-            format="FLAC",
-        )
-        with sf:
-            sf.write(d_signal)
+        if wfdb_archive:
+            with io.BytesIO() as f:
+                with soundfile.SoundFile(
+                    f,
+                    mode="w",
+                    samplerate=96000,
+                    channels=n_sig,
+                    subtype=subtype,
+                    format="FLAC",  # required for file-like
+                ) as sf:
+                    sf.write(d_signal)
+                wfdb_archive.write(os.path.basename(file_name), f.getvalue())
+            return
+        else:
+            sf = soundfile.SoundFile(
+                file_path,
+                mode="w",
+                samplerate=96000,
+                channels=n_sig,
+                subtype=subtype,
+                format="FLAC",
+            )
+            with sf:
+                sf.write(d_signal)
             return
 
     else:
@@ -2549,8 +2572,13 @@ def wr_dat_file(
         b_write = np.append(np.zeros(byte_offset, dtype="uint8"), b_write)
 
     # Write the bytes to the file
-    with open(file_path, "wb") as f:
-        b_write.tofile(f)
+    if wfdb_archive:
+        with io.BytesIO() as f:
+            b_write.tofile(f)
+            wfdb_archive.write(os.path.basename(file_name), f.getvalue())
+    else:
+        with open(file_path, "wb") as f:
+            b_write.tofile(f)
 
 
 def describe_list_indices(full_list):

wfdb/io/archive.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import zipfile
 from contextlib import contextmanager
 
@@ -60,6 +61,23 @@ def close(self):
         if self.zipfile:
             self.zipfile.close()
 
+    def write(self, filename, data):
+        """
+        Write binary data to the archive (replaces if already exists).
+        """
+        # Write to a new temporary archive
+        tmp_path = self.archive_path + ".tmp"
+        with zipfile.ZipFile(self.archive_path, mode="r") as zin:
+            with zipfile.ZipFile(tmp_path, mode="w") as zout:
+                for item in zin.infolist():
+                    if item.filename != filename:
+                        zout.writestr(item, zin.read(item.filename))
+                zout.writestr(filename, data)
+
+        # Replace the original archive
+        shutil.move(tmp_path, self.archive_path)
+        self.zipfile = zipfile.ZipFile(self.archive_path, mode="a")
+
     def create_archive(self, file_list, output_path=None):
         """
         Create a .wfdb archive containing the specified list of files.

wfdb/io/record.py

@@ -905,7 +905,7 @@ def __eq__(self, other, verbose=False):
 
         return True
 
-    def wrsamp(self, expanded=False, write_dir=""):
+    def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None):
         """
         Write a WFDB header file and any associated dat files from this
         object.
@@ -939,7 +939,8 @@ def wrsamp(self, expanded=False, write_dir=""):
         if self.n_sig > 0:
             # Perform signal validity and cohesion checks, and write the
             # associated dat files.
-            self.wr_dats(expanded=expanded, write_dir=write_dir)
+            self.wr_dats(expanded=expanded, write_dir=write_dir, 
+                         wfdb_archive=wfdb_archive)
 
     def _arrange_fields(self, channels, sampfrom, smooth_frames):
         """
@@ -2878,6 +2879,7 @@ def wrsamp(
     base_date=None,
     base_datetime=None,
     write_dir="",
+    archive=False,
 ):
     """
     Write a single segment WFDB record, creating a WFDB header file and any
@@ -2966,6 +2968,7 @@ def wrsamp(
     # Check for valid record name
     if "." in record_name:
         raise Exception("Record name must not contain '.'")
+
     # Check input field combinations
     signal_list = [p_signal, d_signal, e_p_signal, e_d_signal]
     signals_set = sum(1 for var in signal_list if var is not None)
@@ -3064,8 +3067,12 @@ def wrsamp(
     else:
         expanded = False
 
+    wfdb_archive = None
+    if archive:
+        wfdb_archive = get_archive(os.path.join(write_dir, record_name))
+
     # Write the record files - header and associated dat
-    record.wrsamp(write_dir=write_dir, expanded=expanded)
+    record.wrsamp(write_dir=write_dir, expanded=expanded, wfdb_archive=wfdb_archive)
 
 
 def dl_database(