diff --git a/docs/changes.rst b/docs/changes.rst index fc7fc406..eea148df 100644 --- a/docs/changes.rst +++ b/docs/changes.rst @@ -23,6 +23,9 @@ Unreleased **Use Ruff for formatting and linting** The ``dev`` extra now includes ``ruff`` instead of ``black`` and ``pylint``, and the code is formatted with a line length of 88 (the default). +**Read and write .wfdb archive files** + A record can be stored in a single ``.wfdb`` archive file (a ZIP file containing the header and signal files). ``rdrecord``, ``rdsamp``, ``rdheader`` and ``rdann`` read a record or its annotations from an archive when given the path of the archive, and ``wrsamp`` writes a record into an archive with the new ``wfdb_archive`` argument. + Version 4.3.1 (Feb 2026) ----------------------------- diff --git a/docs/io.rst b/docs/io.rst index b47f94bc..fedfb539 100644 --- a/docs/io.rst +++ b/docs/io.rst @@ -34,6 +34,54 @@ WFDB Annotations :members: wrann +WFDB Archives +------------- + +A ``.wfdb`` archive is a ZIP file containing the files of a single +record: its header file, the signal files listed in the header, and +optionally annotation files. The files are at the top level of the +archive, and the header file is named after the archive (for example, +``100.wfdb`` contains ``100.hea``). +Files are stored without compression, so that part of a signal can be +read without decompressing the whole file. Archives made by other tools +may compress their files, and can still be read. + +To read a record from an archive, pass the path of the archive, +including the ``.wfdb`` extension, to ``rdrecord``, ``rdsamp`` or +``rdheader``. Annotations are read the same way with ``rdann``. To +write a record into an archive, pass the path of the archive as the +``wfdb_archive`` argument of ``wrsamp``. + +.. code-block:: python + + import wfdb + + # Write a record into an archive + wfdb.wrsamp("100", fs=360, units=["mV"], sig_name=["MLII"], + p_signal=signal, fmt=["212"], wfdb_archive="100.wfdb") + + # Read the record back, or part of it + record = wfdb.rdrecord("100.wfdb") + record = wfdb.rdrecord("100.wfdb", sampfrom=1000, sampto=2000) + + # Put the files of an existing record, including its annotations, + # into an archive, and read the annotations + from wfdb.io.archive import create_archive + create_archive(["100.hea", "100.dat", "100.atr"], "100.wfdb") + annotation = wfdb.rdann("100.wfdb", "atr") + +An archive holds a single record, so only one header file can be +written to it. Annotation files can be added alongside the record. +Multi-segment records, and archives in cloud storage or on PhysioNet, +are not supported yet. + +.. autofunction:: wfdb.io.archive.create_archive + +.. autoclass:: wfdb.io.archive.WFDBArchive + :members: namelist, exists, file_size, open, header_file_name, + read_header, write, add_file, close, abort + + Cloud Storage Access -------------------- diff --git a/src/wfdb/io/_header.py b/src/wfdb/io/_header.py index 255a2716..662bfe27 100644 --- a/src/wfdb/io/_header.py +++ b/src/wfdb/io/_header.py @@ -290,7 +290,7 @@ def set_defaults(self): for f in sfields: self.set_default(f) - def wrheader(self, write_dir="", expanded=True): + def wrheader(self, write_dir="", expanded=True, wfdb_archive=None): """Write a WFDB header file. The signals are not used. Before writing: @@ -307,6 +307,9 @@ def wrheader(self, write_dir="", expanded=True): Whether the header file should include `samps_per_frame` (this should only be true if the signal files are written using `expanded=True`). + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to write the header file to, instead of + `write_dir`. Returns ------- @@ -338,7 +341,7 @@ def wrheader(self, write_dir="", expanded=True): self.check_field_cohesion(rec_write_fields, list(sig_write_fields)) # Write the header file using the specified fields - self.wr_header_file(rec_write_fields, sig_write_fields, write_dir) + self.wr_header_file(rec_write_fields, sig_write_fields, write_dir, wfdb_archive) def get_write_fields(self): """Get the header fields to write, split into record and signal fields. @@ -514,7 +517,9 @@ def check_field_cohesion(self, rec_write_fields, sig_write_fields): "Each file_name (dat file) specified must have the same byte offset" ) - def wr_header_file(self, rec_write_fields, sig_write_fields, write_dir): + def wr_header_file( + self, rec_write_fields, sig_write_fields, write_dir, wfdb_archive=None + ): """Write a header file using the specified fields. Converts Record attributes into appropriate WFDB format strings. @@ -528,6 +533,9 @@ def wr_header_file(self, rec_write_fields, sig_write_fields, write_dir): being equal to a list of channels to write for each field. write_dir : str The directory in which to write the header file. + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to write the header file to, instead of + `write_dir`. Returns ------- @@ -572,7 +580,11 @@ def wr_header_file(self, rec_write_fields, sig_write_fields, write_dir): comment_lines = ["# " + comment for comment in self.comments] header_lines += comment_lines - util.lines_to_file(self.record_name + ".hea", write_dir, header_lines) + if wfdb_archive is not None: + content = "".join(f"{line}\n" for line in header_lines) + wfdb_archive.write(self.record_name + ".hea", content.encode("utf-8")) + else: + util.lines_to_file(self.record_name + ".hea", write_dir, header_lines) class MultiHeaderMixin(BaseHeaderMixin): diff --git a/src/wfdb/io/_signal.py b/src/wfdb/io/_signal.py index 7bdc5b60..eba0dde2 100644 --- a/src/wfdb/io/_signal.py +++ b/src/wfdb/io/_signal.py @@ -1,3 +1,4 @@ +import io import math import os import posixpath @@ -119,7 +120,7 @@ class SignalMixin: """ - def wr_dats(self, expanded, write_dir): + def wr_dats(self, expanded, write_dir, wfdb_archive=None): """Write all dat files of a record. Use expanded=True to use e_d_signal instead of d_signal. @@ -131,6 +132,9 @@ 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 + An open .wfdb archive to write the dat files to, instead of + `write_dir`. Returns ------- @@ -159,7 +163,9 @@ 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): """Check the cohesion of d_signal/e_d_signal with the other fields. @@ -936,7 +942,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. Parameters @@ -946,6 +952,9 @@ 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 + An open .wfdb archive to write the dat files to, instead of + `write_dir`. Returns ------- @@ -980,6 +989,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 @@ -990,6 +1000,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"): @@ -1094,6 +1105,7 @@ def _rd_segment( no_file=False, sig_data=None, return_res=64, + wfdb_archive=None, ): """Read the digital samples from a single segment record's associated dat file(s). @@ -1141,6 +1153,9 @@ def _rd_segment( 32, 16, and 8, where the value represents the numpy int or float dtype. Note that the value cannot be 8 when physical is True since there is no float8 format. + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to read the signal files from instead of + `dir_name` or `pn_dir`. Returns ------- @@ -1238,6 +1253,7 @@ def _rd_segment( sampto=sampto, no_file=no_file, sig_data=sig_data, + wfdb_archive=wfdb_archive, ) # Copy over the wanted signals @@ -1262,6 +1278,7 @@ def _rd_dat_signals( sampto, no_file=False, sig_data=None, + wfdb_archive=None, ): """Read all signals from a WFDB dat file. @@ -1299,6 +1316,9 @@ def _rd_dat_signals( sig_data : ndarray, optional The signal data that would normally be imported using the associated .dat and .hea files. Should only be used when no_file is set to True. + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to read the signal files from instead of + `dir_name` or `pn_dir`. Returns ------- @@ -1360,10 +1380,17 @@ def _rd_dat_signals( samps_per_frame=samps_per_frame, start_frame=sampfrom, end_frame=sampto, + wfdb_archive=wfdb_archive, ) else: data_to_read = _rd_dat_file( - file_name, dir_name, pn_dir, fmt, start_byte, n_read_samples + file_name, + dir_name, + pn_dir, + fmt, + start_byte, + n_read_samples, + wfdb_archive=wfdb_archive, ) if extra_flat_samples: @@ -1601,7 +1628,9 @@ def _required_byte_num(mode, fmt, n_samp): return int(n_bytes) -def _rd_dat_file(file_name, dir_name, pn_dir, fmt, start_byte, n_samp): +def _rd_dat_file( + file_name, dir_name, pn_dir, fmt, start_byte, n_samp, wfdb_archive=None +): """Read data from a dat file, either local or remote, into a 1d numpy array. This is the lowest level dat reading function (along with @@ -1626,6 +1655,9 @@ def _rd_dat_file(file_name, dir_name, pn_dir, fmt, start_byte, n_samp): The total number of samples to read. Does NOT need to create whole blocks for special format. Any number of samples should be readable. + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to read the signal files from instead of + `dir_name` or `pn_dir`. Returns ------- @@ -1657,8 +1689,16 @@ def _rd_dat_file(file_name, dir_name, pn_dir, fmt, start_byte, n_samp): element_count = n_samp byte_count = n_samp * BYTES_PER_SAMPLE[fmt] + # Dat file in a .wfdb archive + if wfdb_archive is not None: + with wfdb_archive.open(file_name) as fp: + fp.seek(start_byte) + sig_data = util.fromfile( + fp, dtype=np.dtype(DATA_LOAD_TYPES[fmt]), count=element_count + ) + # Local or cloud dat file - if pn_dir is None: + elif pn_dir is None: with fsspec.open(os.path.join(dir_name, file_name), "rb") as fp: fp.seek(start_byte) sig_data = util.fromfile( @@ -1806,6 +1846,7 @@ def _rd_compressed_file( samps_per_frame, start_frame, end_frame, + wfdb_archive=None, ): """Read data from a compressed file into a 1D numpy array. @@ -1832,6 +1873,9 @@ def _rd_compressed_file( The starting frame number to read. end_frame : int The ending frame number to read. + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to read the signal files from instead of + `dir_name` or `pn_dir`. Returns ------- @@ -1856,10 +1900,14 @@ def _rd_compressed_file( "sampling rate and samples per frame" ) - if pn_dir is None: - file_name = os.path.join(dir_name, file_name) + if wfdb_archive is not None: + open_file = wfdb_archive.open(file_name) + else: + if pn_dir is None: + file_name = os.path.join(dir_name, file_name) + open_file = _coreio._open_file(pn_dir, file_name, "rb") - with _coreio._open_file(pn_dir, file_name, "rb") as fp: + with open_file as fp: signature = fp.read(4) if signature != b"fLaC": raise ValueError(f"{fp.name} is not a FLAC file") @@ -2267,6 +2315,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 endianness issues. @@ -2290,6 +2339,9 @@ def wr_dat_file( The samples/frame for each signal of the dat file. write_dir : str, optional The directory to write the output file to. + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to write the dat file to, instead of + `write_dir`. Returns ------- @@ -2461,8 +2513,9 @@ def wr_dat_file( else: raise ValueError(f"unknown format ({fmt})") + target = io.BytesIO() if wfdb_archive is not None else file_path sf = soundfile.SoundFile( - file_path, + target, mode="w", samplerate=96000, channels=n_sig, @@ -2471,7 +2524,9 @@ def wr_dat_file( ) with sf: sf.write(d_signal) - return + if wfdb_archive is not None: + wfdb_archive.write(file_name, target.getbuffer()) + return else: raise ValueError( @@ -2491,8 +2546,11 @@ 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 is not None: + wfdb_archive.write(file_name, np.ascontiguousarray(b_write)) + else: + with open(file_path, "wb") as f: + b_write.tofile(f) def describe_list_indices(full_list): @@ -2529,7 +2587,13 @@ def describe_list_indices(full_list): def _infer_sig_len( - file_name, fmt, tsamps_per_frame, byte_offset, dir_name, pn_dir=None + file_name, + fmt, + tsamps_per_frame, + byte_offset, + dir_name, + pn_dir=None, + wfdb_archive=None, ): """Infer the length of a signal from a dat file. @@ -2549,6 +2613,9 @@ def _infer_sig_len( pn_dir : str, optional The PhysioNet directory where the dat file(s) are located, if the dat file(s) are on a PhysioNet server. + wfdb_archive : WFDBArchive, optional + An open .wfdb archive to read the dat file from instead of + `dir_name` or `pn_dir`. Returns ------- @@ -2562,9 +2629,13 @@ def _infer_sig_len( """ from wfdb.io.record import CLOUD_PROTOCOLS + # If the dat file is in a .wfdb archive, get its uncompressed size + if wfdb_archive is not None: + file_size = wfdb_archive.file_size(file_name) + # If this is a cloud path, use posixpath to construct the path and fsspec to open # file - if any(dir_name.startswith(proto) for proto in CLOUD_PROTOCOLS): + elif any(dir_name.startswith(proto) for proto in CLOUD_PROTOCOLS): with fsspec.open(posixpath.join(dir_name, file_name), mode="rb") as f: file_size = f.seek(0, os.SEEK_END) diff --git a/src/wfdb/io/annotation.py b/src/wfdb/io/annotation.py index 20f6b7b1..b171dda2 100644 --- a/src/wfdb/io/annotation.py +++ b/src/wfdb/io/annotation.py @@ -10,6 +10,7 @@ from wfdb.io import _header, download, record, util from wfdb.io._coreio import CLOUD_PROTOCOLS +from wfdb.io.archive import WFDBArchive, is_archive_path class Annotation: @@ -1872,7 +1873,10 @@ def rdann( ---------- record_name : str The record name of the WFDB annotation file. ie. for file '100.atr', - record_name='100'. The path to the file can be a cloud URL. + record_name='100'. The path to the file can be a cloud URL. To read + the annotation file from a local .wfdb archive, pass the path of + the archive including the .wfdb extension (for example + '100.wfdb' to read '100.atr' from the archive). extension : str The annotatator extension of the annotation file. ie. for file '100.atr', extension='atr'. @@ -1914,8 +1918,13 @@ def rdann( Examples -------- >>> ann = wfdb.rdann('sample-data/100', 'atr', sampto=300000) + >>> ann = wfdb.rdann('100.wfdb', 'atr') """ + is_archive = is_archive_path(record_name) + if is_archive: + record._check_archive_read_args(record_name, pn_dir) + if pn_dir is not None: # check to make sure a cloud path isn't being passed under pn_dir if any(pn_dir.startswith(proto) for proto in CLOUD_PROTOCOLS): @@ -1931,7 +1940,11 @@ def rdann( return_label_elements = check_read_inputs(sampfrom, sampto, return_label_elements) # Read the file in byte pairs - filebytes = load_byte_pairs(record_name, extension, pn_dir) + if is_archive: + filebytes, ann_record_name = _load_archive_byte_pairs(record_name, extension) + else: + filebytes = load_byte_pairs(record_name, extension, pn_dir) + ann_record_name = os.path.split(record_name)[1] # Get WFDB annotation fields from the file bytes sample, label_store, subtype, chan, num, aux_note = proc_ann_bytes( @@ -1971,7 +1984,7 @@ def rdann( # Create the annotation object annotation = Annotation( - record_name=os.path.split(record_name)[1], + record_name=ann_record_name, extension=extension, sample=sample, label_store=label_store, @@ -2080,6 +2093,33 @@ def load_byte_pairs(record_name, extension, pn_dir): return filebytes +def _load_archive_byte_pairs(archive_path, extension): + """Load an annotation file from a .wfdb archive as pairs of bytes. + + The annotation file is named after the record's header file in the archive. + + Parameters + ---------- + archive_path : str + The path of the .wfdb archive. + extension : str + The annotator extension of the annotation file. + + Returns + ------- + filebytes : ndarray + The annotation file converted to an Nx2 array of unsigned bytes. + record_name : str + The name of the record that the annotation file belongs to. + + """ + with WFDBArchive(archive_path, mode="r") as archive: + record_name = archive.header_file_name()[: -len(".hea")] + with archive.open(f"{record_name}.{extension}") as f: + filebytes = util.fromfile(f, ">> with WFDBArchive("record.wfdb", mode="w") as archive: + ... wfdb.wrsamp("record", fs=250, units=["mV"], sig_name=["I"], + ... p_signal=signal, fmt=["16"], wfdb_archive=archive) + + """ + + def __init__(self, path, mode="r"): + if mode not in ("r", "w"): + raise ValueError("mode must be 'r' or 'w'") + + path = os.fspath(path) + if not is_archive_path(path): + path += ARCHIVE_EXTENSION + self.path = path + self.mode = mode + self._tmp_path = None + + if mode == "r": + if not os.path.isfile(path): + raise FileNotFoundError(f"Archive not found: {path}") + try: + self._zipfile = zipfile.ZipFile(path, mode="r") + except zipfile.BadZipFile as e: + raise ValueError(f"Not a valid WFDB archive: {path}") from e + else: + self._tmp_path = _create_temp_file(path) + try: + # Files dated before 1980 (which ZIP can't represent) get the + # earliest possible date instead of raising an error + self._zipfile = zipfile.ZipFile( + self._tmp_path, mode="w", strict_timestamps=False + ) + except BaseException: + os.remove(self._tmp_path) + raise + + # Close the archive, and remove an unsaved temporary file, if the + # object is discarded without being closed + self._finalizer = weakref.finalize( + self, _discard, self._zipfile, self._tmp_path + ) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is None: + self.close() + else: + self.abort() + + @property + def record_name(self): + """The record name implied by the archive file name.""" + return os.path.basename(self.path)[: -len(ARCHIVE_EXTENSION)] + + def namelist(self): + """Return the names of the files in the archive.""" + self._require_open() + return self._zipfile.namelist() + + def exists(self, file_name): + """Check whether a file is in the archive.""" + self._require_open() + try: + self._zipfile.getinfo(file_name) + except KeyError: + return False + return True + + def file_size(self, file_name): + """Return the uncompressed size of a file in the archive, in bytes.""" + return self._getinfo(file_name).file_size + + def open(self, file_name): + """Open a file in the archive for reading as a binary file object.""" + self._require_mode("r") + return self._zipfile.open(self._getinfo(file_name), mode="r") + + def header_file_name(self): + """Find the header file of the record in the archive. + + The header is expected to be named after the archive (for + example 'record.hea' in 'record.wfdb'). If there is no such + file but the archive contains exactly one header file, that + file is used. + + Returns + ------- + str + The name of the header file in the archive. + + """ + file_name = self.record_name + ".hea" + if not self.exists(file_name): + headers = self._header_names() + if len(headers) != 1: + raise FileNotFoundError( + f"Could not find '{file_name}' in archive '{self.path}'" + + ( + f" (it contains several header files: {headers})" + if headers + else "" + ) + ) + file_name = headers[0] + return file_name + + def read_header(self): + """Return the content of the record's header file. + + See `header_file_name` for how the header file is found. + + Returns + ------- + str + The content of the header file. + + """ + file_name = self.header_file_name() + with self.open(file_name) as f: + # Read at most one byte more than allowed, whatever size the archive + # declares, so that a damaged or malicious header isn't read in full + content = f.read(MAX_HEADER_SIZE + 1) + if len(content) > MAX_HEADER_SIZE: + raise ValueError( + f"Header file '{file_name}' in archive '{self.path}' is " + f"larger than {MAX_HEADER_SIZE} bytes" + ) + return content.decode("ascii", errors="ignore") + + def write(self, file_name, data): + """Add a file to an archive opened for writing. + + The data is streamed into the archive, so a large signal is not copied in + memory. + + Parameters + ---------- + file_name : str + The name of the file in the archive. + data : bytes-like object + The content of the file, such as `bytes`, or a `memoryview` of a + contiguous array. + + """ + self._check_new_file(file_name) + data = memoryview(data).cast("B") + # The same file information that ZipFile.writestr uses + info = zipfile.ZipInfo(file_name, date_time=time.localtime()[:6]) + info.external_attr = 0o600 << 16 + # ZIP64 is needed for files of 2 GiB or more, and has to be chosen up front + with self._zipfile.open( + info, mode="w", force_zip64=data.nbytes >= zipfile.ZIP64_LIMIT + ) as f: + f.write(data) + + def add_file(self, path): + """Copy a file from disk into an archive opened for writing. + + Parameters + ---------- + path : str or os.PathLike + The path of the file. It is stored under its base name. + + """ + file_name = os.path.basename(path) + self._check_new_file(file_name) + self._zipfile.write(path, arcname=file_name) + + def close(self): + """Close the archive. + + In write mode, this moves the new archive into place. + """ + if self._zipfile is None: + return + if self.mode == "w": + try: + self._check_header_name() + except BaseException: + self.abort() + raise + self._finalizer.detach() + try: + self._zipfile.close() + if self._tmp_path is not None: + # A new archive has the default permissions for new files. A + # replaced archive keeps its permissions. + try: + mode = stat.S_IMODE(os.stat(self.path).st_mode) + except FileNotFoundError: + pass + else: + os.chmod(self._tmp_path, mode) + os.replace(self._tmp_path, self.path) + except BaseException: + _discard(None, self._tmp_path) + raise + finally: + self._zipfile = None + self._tmp_path = None + + def abort(self): + """Close the archive without saving it. + + In read mode, this is the same as `close`. + """ + if self._zipfile is None: + return + self._finalizer() + self._zipfile = None + self._tmp_path = None + + def _getinfo(self, file_name): + self._require_open() + try: + return self._zipfile.getinfo(file_name) + except KeyError: + raise FileNotFoundError( + f"Could not find '{file_name}' in archive '{self.path}'" + ) from None + + def _check_new_file(self, file_name): + """Check that a file can be added to an archive opened for writing.""" + self._require_mode("w") + if not _FILE_NAME.fullmatch(file_name): + raise ValueError( + f"Cannot add '{file_name}' to '{self.path}': files in an archive " + "are at the top level, and their names are made of ASCII letters, " + "digits, hyphens and underscores, with an extension" + ) + if self.exists(file_name): + raise ValueError(f"'{file_name}' was already written to the archive") + if _HEADER_NAME.fullmatch(file_name): + headers = self._header_names() + if headers: + raise ValueError( + f"Cannot add '{file_name}' to '{self.path}': an archive " + f"holds a single record, and it already contains the " + f"header '{headers[0]}'" + ) + + def _check_header_name(self): + """Check that the record's header file is named after the archive.""" + expected = self.record_name + ".hea" + headers = self._header_names() + if headers and headers != [expected]: + raise ValueError( + f"Cannot save '{self.path}': the header file of the record in an " + f"archive must be named after the archive ('{expected}'), but it " + f"is '{headers[0]}'" + ) + + def _header_names(self): + """Return the names of the header files in the archive. + + Only files at the top level that are named after a valid record name count, + so metadata files that macOS adds, such as '__MACOSX/._100.hea' (added by + Finder when it zips '100.hea') or '._100.hea' (written next to '100.hea' on + some drives), are ignored. + """ + return [n for n in self.namelist() if _HEADER_NAME.fullmatch(n)] + + def _require_open(self): + if self._zipfile is None: + raise ValueError(f"Archive '{self.path}' is closed") + + def _require_mode(self, mode): + self._require_open() + if self.mode != mode: + action = "read from" if mode == "r" else "written to" + raise ValueError( + f"Archive '{self.path}' was opened in mode '{self.mode}' " + f"and cannot be {action}" + ) + + +def _discard(zip_file, tmp_path): + """Close a zip file and remove its temporary file, if there is one. + + The temporary file is removed even if closing the zip file fails. + """ + try: + if zip_file is not None: + zip_file.close() + finally: + if tmp_path is not None: + with contextlib.suppress(FileNotFoundError): + os.remove(tmp_path) + + +def _create_temp_file(path): + """Create an empty temporary file next to `path`, and return its path. + + The file is named after `path`, with a random part. It is created with the default + permissions for new files (the umask is applied by the operating system), unlike + `tempfile.mkstemp`, which makes files readable only by their owner. + """ + directory = os.path.dirname(os.path.abspath(path)) + os.makedirs(directory, exist_ok=True) + for _ in range(100): + tmp_path = os.path.join( + directory, f".{os.path.basename(path)}.{secrets.token_hex(4)}.tmp" + ) + try: + fd = os.open(tmp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o666) + except FileExistsError: + continue + os.close(fd) + return tmp_path + raise FileExistsError(f"Could not create a temporary file next to '{path}'") + + +def create_archive(file_paths, archive_path): + """Create a .wfdb archive from the files of a record on disk. + + Parameters + ---------- + file_paths : list of str + The paths of the files to add: the header file and its signal + files. Each file is stored under its base name. + archive_path : str or os.PathLike + The path of the archive to create, which must be named after the + record. An existing archive at this path is replaced. + + Examples + -------- + >>> create_archive(["100.hea", "100.dat"], "100.wfdb") + + """ + with WFDBArchive(archive_path, mode="w") as archive: + for path in file_paths: + archive.add_file(path) + + +def is_archive_path(record_name): + """Check whether a record name refers to a .wfdb archive. + + The extension is not case-sensitive, so 'record.WFDB' is an archive too. + """ + return os.fspath(record_name).lower().endswith(ARCHIVE_EXTENSION) + + +@contextmanager +def archive_for_writing(wfdb_archive, write_dir=""): + """Get a WFDBArchive to write a record to. + + Parameters + ---------- + wfdb_archive : str, os.PathLike, WFDBArchive or None + A path opens a new archive that is saved when the context exits, + or discarded if an error occurs. A WFDBArchive opened in write + mode is used as it is and left open. None gives None. + write_dir : str, optional + The directory that a relative path is relative to, as for the + files of a record written without an archive. + + """ + if wfdb_archive is None: + yield None + elif isinstance(wfdb_archive, WFDBArchive): + wfdb_archive._require_mode("w") + yield wfdb_archive + elif isinstance(wfdb_archive, (str, os.PathLike)): + path = os.path.join(write_dir, wfdb_archive) + with WFDBArchive(path, mode="w") as archive: + yield archive + else: + raise TypeError( + "wfdb_archive must be a path or a WFDBArchive opened in write mode" + ) diff --git a/src/wfdb/io/record.py b/src/wfdb/io/record.py index 107b0fc9..52a879b3 100644 --- a/src/wfdb/io/record.py +++ b/src/wfdb/io/record.py @@ -1,3 +1,4 @@ +import contextlib import datetime import multiprocessing.dummy import os @@ -10,6 +11,7 @@ from wfdb.io import _header, _signal, _url, download, header, util from wfdb.io._coreio import CLOUD_PROTOCOLS +from wfdb.io.archive import WFDBArchive, archive_for_writing, is_archive_path # -------------- WFDB Signal Calibration and Classification ---------- # @@ -870,7 +872,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. Parameters @@ -880,6 +882,12 @@ def wrsamp(self, expanded=False, write_dir=""): of the uniform signal (d_signal). write_dir : str, optional The directory in which to write the files. + wfdb_archive : str, os.PathLike or WFDBArchive, optional + Write the record into a .wfdb archive instead of writing separate + files to `write_dir`. A path creates a new archive, replacing any + existing archive at that path. A relative path is relative to + `write_dir`. A WFDBArchive must be opened in write mode, and is left + open. Returns ------- @@ -896,13 +904,16 @@ def wrsamp(self, expanded=False, write_dir=""): checksums[ch] = old_val self.checksum = checksums - # Perform field validity and cohesion checks, and write the - # header file. - self.wrheader(write_dir=write_dir, expanded=expanded) - 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) + with archive_for_writing(wfdb_archive, write_dir) as archive: + # Perform field validity and cohesion checks, and write the + # header file. + self.wrheader(write_dir=write_dir, expanded=expanded, wfdb_archive=archive) + 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, wfdb_archive=archive + ) def _arrange_fields(self, channels, sampfrom, smooth_frames): """Arrange/edit object fields to reflect user channel and/or signal range input. @@ -1759,7 +1770,9 @@ def rdheader(record_name, pn_dir=None, rd_segments=False): supported cloud URI, they are read directly from cloud storage via `fsspec`. Supported URIs: `s3://`, `gs://`, `az://`, `azureml://`. Install the appropriate `fsspec` backend - (e.g., `s3fs`, `gcsfs`, `adlfs`). + (e.g., `s3fs`, `gcsfs`, `adlfs`). To read the header from a + local .wfdb archive, pass the path of the archive including the + .wfdb extension. pn_dir : str, optional Option used to stream data from Physionet. The Physionet database directory from which to find the required record files. @@ -1782,10 +1795,17 @@ def rdheader(record_name, pn_dir=None, rd_segments=False): """ dir_name, base_record_name = os.path.split(record_name) file_name = f"{base_record_name}.hea" + is_archive = is_archive_path(record_name) + + # Read the header from a .wfdb archive + if is_archive: + _check_archive_read_args(record_name, pn_dir) + with WFDBArchive(record_name, mode="r") as archive: + header_content = archive.read_header() # If this is a cloud path, use posixpath to construct the path and fsspec to open # file - if any(dir_name.startswith(proto) for proto in CLOUD_PROTOCOLS): + elif any(dir_name.startswith(proto) for proto in CLOUD_PROTOCOLS): with fsspec.open(posixpath.join(dir_name, file_name), mode="r") as f: header_content = f.read() @@ -1817,6 +1837,45 @@ def rdheader(record_name, pn_dir=None, rd_segments=False): ) as f: header_content = f.read() + record = _parse_header_content(header_content) + + # If specified, read the segment headers of a multi-segment record + if rd_segments and isinstance(record, MultiRecord): + if is_archive: + raise NotImplementedError( + "Multi-segment records in .wfdb archives are not supported" + ) + record.segments = [] + # Get the base record name (could be empty) + for s in record.seg_name: + if s == "~": + record.segments.append(None) + else: + record.segments.append(rdheader(os.path.join(dir_name, s), pn_dir)) + # Fill in the sig_name attribute + record.sig_name = record.get_sig_name() + # Fill in the sig_segments attribute + record.sig_segments = record.get_sig_segments() + + return record + + +def _parse_header_content(header_content): + """Create a `Record` or `MultiRecord` object from the content of a header file. + + The segment headers of a multi-segment record are not read. + + Parameters + ---------- + header_content : str + The content of the header file. + + Returns + ------- + record : Record or MultiRecord + The record described by the header. + + """ # Separate comment and non-comment lines header_lines, comment_lines = header.parse_header_content(header_content) @@ -1860,26 +1919,22 @@ def rdheader(record_name, pn_dir=None, rd_segments=False): else: record.layout = "fixed" - # If specified, read the segment headers - if rd_segments: - record.segments = [] - # Get the base record name (could be empty) - for s in record.seg_name: - if s == "~": - record.segments.append(None) - else: - record.segments.append(rdheader(os.path.join(dir_name, s), pn_dir)) - # Fill in the sig_name attribute - record.sig_name = record.get_sig_name() - # Fill in the sig_segments attribute - record.sig_segments = record.get_sig_segments() - # Set the comments field record.comments = [line.strip(" \t#") for line in comment_lines] return record +def _check_archive_read_args(record_name, pn_dir): + """Check that a .wfdb archive is read from a supported location.""" + if pn_dir is not None: + raise ValueError("pn_dir cannot be used to read a .wfdb archive") + if any(os.fspath(record_name).startswith(p) for p in CLOUD_PROTOCOLS): + raise NotImplementedError( + "Reading .wfdb archives from cloud storage is not supported" + ) + + def rdrecord( record_name, sampfrom=0, @@ -1912,7 +1967,9 @@ def rdrecord( supported cloud URI, they are read directly from cloud storage via `fsspec`. Supported URIs: `s3://`, `gs://`, `az://`, `azureml://`. Install the appropriate `fsspec` backend - (e.g., `s3fs`, `gcsfs`, `adlfs`). + (e.g., `s3fs`, `gcsfs`, `adlfs`). To read a record from a local + .wfdb archive, pass the path of the archive including the .wfdb + extension. Multi-segment records are not supported in archives. sampfrom : int, optional The starting sample number to read for all channels. sampto : int, 'end', optional @@ -1987,6 +2044,56 @@ def rdrecord( -------- >>> record = wfdb.rdrecord('sample-data/test01_00s', sampfrom=800, channels=[1, 3]) + >>> record = wfdb.rdrecord('100.wfdb') + + """ + if is_archive_path(record_name): + _check_archive_read_args(record_name, pn_dir) + archive = WFDBArchive(record_name, mode="r") + else: + archive = contextlib.nullcontext() + + with archive as wfdb_archive: + return _rdrecord( + record_name=record_name, + sampfrom=sampfrom, + sampto=sampto, + channels=channels, + physical=physical, + pn_dir=pn_dir, + m2s=m2s, + smooth_frames=smooth_frames, + ignore_skew=ignore_skew, + return_res=return_res, + force_channels=force_channels, + channel_names=channel_names, + warn_empty=warn_empty, + wfdb_archive=wfdb_archive, + ) + + +def _rdrecord( + record_name, + sampfrom, + sampto, + channels, + physical, + pn_dir, + m2s, + smooth_frames, + ignore_skew, + return_res, + force_channels, + channel_names, + warn_empty, + wfdb_archive, +): + """Read a WFDB record. + + See `rdrecord` for the parameters. + + If `wfdb_archive` is set, it is the open WFDBArchive that + `record_name` refers to, and the signal files are read from it. """ dir_name, base_record_name = os.path.split(record_name) @@ -2007,7 +2114,14 @@ def rdrecord( dir_list[0], download.get_version(dir_list[0]), *dir_list[1:] ) - record = rdheader(record_name, pn_dir=pn_dir, rd_segments=False) + if wfdb_archive is not None: + record = _parse_header_content(wfdb_archive.read_header()) + if isinstance(record, MultiRecord): + raise NotImplementedError( + "Multi-segment records in .wfdb archives are not supported" + ) + else: + record = rdheader(record_name, pn_dir=pn_dir, rd_segments=False) # Set defaults for sampto and channels input variables if sampto is None: @@ -2033,6 +2147,7 @@ def rdrecord( byte_offset=record.byte_offset[0], dir_name=dir_name, pn_dir=pn_dir, + wfdb_archive=wfdb_archive, ) sampto = record.sig_len @@ -2109,6 +2224,7 @@ def rdrecord( no_file=no_file, sig_data=sig_data, return_res=return_res, + wfdb_archive=wfdb_archive, ) # Only 1 sample/frame, or frames are smoothed. Return uniform numpy array @@ -2830,6 +2946,7 @@ def wrsamp( base_date=None, base_datetime=None, write_dir="", + wfdb_archive=None, ): """Write a single segment WFDB record with a header and any dat files. @@ -2885,6 +3002,12 @@ def wrsamp( setting both `base_date` and `base_time`. write_dir : str, optional The directory in which to write the files. + wfdb_archive : str, os.PathLike or WFDBArchive, optional + Write the record into a .wfdb archive instead of writing separate + files to `write_dir`. A path creates a new archive, replacing any + existing archive at that path. A relative path is relative to + `write_dir`. A WFDBArchive must be opened in write mode, and is left + open. Returns ------- @@ -2909,6 +3032,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 the same record into a .wfdb archive + >>> 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 @@ -3013,7 +3140,7 @@ def wrsamp( expanded = False # 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( diff --git a/tests/test_archive.py b/tests/test_archive.py new file mode 100644 index 00000000..d15b485b --- /dev/null +++ b/tests/test_archive.py @@ -0,0 +1,742 @@ +import gc +import os +import shutil +import zipfile + +import numpy as np +import pytest + +import wfdb +from wfdb.io import archive as wfdb_archive_module +from wfdb.io.archive import WFDBArchive, create_archive + + +def record_files(record_name): + """Return the header and signal file paths of a record on disk.""" + header = wfdb.rdheader(record_name) + dir_name = os.path.dirname(record_name) + files = [record_name + ".hea"] + for file_name in dict.fromkeys(header.file_name or []): + files.append(os.path.join(dir_name, file_name)) + return files + + +def zip_files(file_paths, archive_path): + """Put files into a zip file the way other software might, without any checks.""" + with zipfile.ZipFile(archive_path, "w") as zf: + for path in file_paths: + zf.write(path, arcname=os.path.basename(path)) + + +def archive_sample_record(record_name, tmp_path): + """Put a record from sample-data into a .wfdb archive and return its path.""" + archive_path = os.fspath(tmp_path / (os.path.basename(record_name) + ".wfdb")) + create_archive(record_files(record_name), archive_path) + return archive_path + + +def assert_records_equal(record, expected): + assert record.n_sig == expected.n_sig + assert record.sig_len == expected.sig_len + assert record.sig_name == expected.sig_name + assert record.fmt == expected.fmt + for attr in ["d_signal", "p_signal"]: + if getattr(expected, attr) is not None: + np.testing.assert_array_equal( + getattr(record, attr), getattr(expected, attr) + ) + for attr in ["e_d_signal", "e_p_signal"]: + if getattr(expected, attr) is not None: + for got, want in zip(getattr(record, attr), getattr(expected, attr)): + np.testing.assert_array_equal(got, want) + + +@pytest.mark.parametrize( + "record_name", + [ + "sample-data/100", + "sample-data/test01_00s", + "sample-data/flacformats", + "sample-data/a103l", + ], +) +@pytest.mark.parametrize("physical", [True, False]) +def test_read_archive_matches_files(record_name, physical, tmp_path): + """Reading from an archive gives the same result as reading the record files.""" + archive_path = archive_sample_record(record_name, tmp_path) + + record = wfdb.rdrecord(archive_path, physical=physical) + expected = wfdb.rdrecord(record_name, physical=physical) + + assert_records_equal(record, expected) + + +def test_read_archive_expanded(tmp_path): + """Signals with several samples per frame can be read without smoothing.""" + record_name = "sample-data/test01_00s_frame" + archive_path = archive_sample_record(record_name, tmp_path) + + record = wfdb.rdrecord(archive_path, smooth_frames=False) + expected = wfdb.rdrecord(record_name, smooth_frames=False) + + assert_records_equal(record, expected) + + +def test_read_archive_channels_and_range(tmp_path): + archive_path = archive_sample_record("sample-data/100", tmp_path) + + record = wfdb.rdrecord(archive_path, sampfrom=1000, sampto=3000, channels=[1]) + expected = wfdb.rdrecord( + "sample-data/100", sampfrom=1000, sampto=3000, channels=[1] + ) + + assert_records_equal(record, expected) + + +def test_rdsamp_and_rdheader_archive(tmp_path): + archive_path = archive_sample_record("sample-data/100", tmp_path) + + signal, fields = wfdb.rdsamp(archive_path, sampto=500) + expected_signal, expected_fields = wfdb.rdsamp("sample-data/100", sampto=500) + np.testing.assert_array_equal(signal, expected_signal) + assert fields == expected_fields + + header = wfdb.rdheader(archive_path) + assert header.record_name == "100" + assert header.sig_name == expected_fields["sig_name"] + + +def test_read_archive_without_sig_len(tmp_path): + """The signal length is inferred from the dat file if the header lacks it.""" + expected = wfdb.rdrecord("sample-data/100") + with open("sample-data/100.hea") as f: + lines = f.read().splitlines() + # Drop sig_len (and anything after it) from the record line + lines[0] = " ".join(lines[0].split()[:3]) + archive_path = os.fspath(tmp_path / "100.wfdb") + with WFDBArchive(archive_path, mode="w") as archive: + archive.write("100.hea", ("\n".join(lines) + "\n").encode()) + archive.add_file("sample-data/100.dat") + + record = wfdb.rdrecord(archive_path) + + assert record.sig_len == expected.sig_len + np.testing.assert_array_equal(record.p_signal, expected.p_signal) + + +def test_read_archive_twice(tmp_path): + archive_path = archive_sample_record("sample-data/100", tmp_path) + + first = wfdb.rdrecord(archive_path, sampto=100) + second = wfdb.rdrecord(archive_path, sampto=100) + + np.testing.assert_array_equal(first.p_signal, second.p_signal) + + +def test_read_replaced_archive(tmp_path): + """Replacing an archive and reading it again returns the new record.""" + archive_path = archive_sample_record("sample-data/100", tmp_path) + assert wfdb.rdrecord(archive_path).sig_name == ["MLII", "V5"] + + wfdb.wrsamp("100", wfdb_archive=archive_path, **write_kwargs()) + + assert wfdb.rdrecord(archive_path).sig_name == ["I", "II"] + + +def test_read_archive_header_named_differently(tmp_path): + """A single header can be read even if it isn't named after the archive. + + wfdb doesn't write such archives, but other software might, or an archive might be + renamed. + """ + archive_path = os.fspath(tmp_path / "renamed.wfdb") + zip_files(record_files("sample-data/100"), archive_path) + + record = wfdb.rdrecord(archive_path, sampto=100) + + assert record.record_name == "100" + + +def test_read_archive_several_headers(tmp_path): + """An archive with several headers, none named after it, can't be read. + + Such an archive can only be made by other software, and it is not clear which record + to read. + """ + archive_path = os.fspath(tmp_path / "both.wfdb") + with zipfile.ZipFile(archive_path, "w") as zf: + for path in record_files("sample-data/100") + ["sample-data/test01_00s.hea"]: + zf.write(path, arcname=os.path.basename(path)) + + with pytest.raises(FileNotFoundError, match="several header files"): + wfdb.rdrecord(archive_path) + + +@pytest.mark.parametrize("prefix", ["__MACOSX/._", "._"]) +def test_read_archive_ignores_macos_metadata(tmp_path, prefix): + """Metadata files that macOS adds next to the record's files are ignored. + + Finder adds '__MACOSX/._100.hea' when it zips '100.hea', and macOS writes + '._100.hea' next to '100.hea' on some drives. Neither counts as a second header, + as only files at the top level named after a valid record name do. + """ + archive_path = os.fspath(tmp_path / "renamed.wfdb") + with zipfile.ZipFile(archive_path, "w") as zf: + for path in record_files("sample-data/100"): + file_name = os.path.basename(path) + zf.write(path, arcname=file_name) + zf.writestr(prefix + file_name, b"\x00\x05\x16\x07") + + record = wfdb.rdrecord(archive_path, sampto=100) + + assert record.record_name == "100" + + +def test_read_compressed_archive(tmp_path): + """An archive whose files are compressed, as other tools make them, can be read.""" + archive_path = os.fspath(tmp_path / "100.wfdb") + with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: + for path in record_files("sample-data/100") + ["sample-data/100.atr"]: + zf.write(path, arcname=os.path.basename(path)) + + record = wfdb.rdrecord(archive_path, sampfrom=1000, sampto=2000) + archive_annotation = wfdb.rdann(archive_path, "atr") + + expected = wfdb.rdrecord("sample-data/100", sampfrom=1000, sampto=2000) + assert_records_equal(record, expected) + np.testing.assert_array_equal( + archive_annotation.sample, wfdb.rdann("sample-data/100", "atr").sample + ) + + +def test_read_archive_missing_dat_file(tmp_path): + """A dat file named in the header but missing from the archive is an error. + + This is the case even if the archive contains another dat file. + """ + with open("sample-data/100.hea") as f: + header = f.read().replace("100.dat", "other.dat") + archive_path = os.fspath(tmp_path / "100.wfdb") + with WFDBArchive(archive_path, mode="w") as archive: + archive.write("100.hea", header.encode()) + archive.add_file("sample-data/100.dat") + + with pytest.raises(FileNotFoundError, match="other.dat"): + wfdb.rdrecord(archive_path) + + +def test_read_archive_without_header(tmp_path): + archive_path = os.fspath(tmp_path / "100.wfdb") + create_archive(["sample-data/100.dat"], archive_path) + + with pytest.raises(FileNotFoundError, match="100.hea"): + wfdb.rdrecord(archive_path) + + +def test_read_missing_archive(tmp_path): + with pytest.raises(FileNotFoundError, match="Archive not found"): + wfdb.rdrecord(os.fspath(tmp_path / "missing.wfdb")) + + +def test_read_invalid_archive(tmp_path): + archive_path = tmp_path / "invalid.wfdb" + archive_path.write_bytes(b"not a zip file") + + with pytest.raises(ValueError, match="Not a valid WFDB archive"): + wfdb.rdrecord(os.fspath(archive_path)) + + +def test_read_archive_header_too_large(tmp_path, monkeypatch): + archive_path = archive_sample_record("sample-data/100", tmp_path) + monkeypatch.setattr(wfdb_archive_module, "MAX_HEADER_SIZE", 10) + + with pytest.raises(ValueError, match="larger than 10 bytes"): + wfdb.rdrecord(archive_path) + + +def test_read_multi_segment_archive_not_supported(tmp_path): + record_dir = "sample-data/multi-segment/fixed1" + archive_path = os.fspath(tmp_path / "v102s.wfdb") + # create_archive only accepts one header, so build the archive the way + # other software might + with zipfile.ZipFile(archive_path, "w") as zf: + for name in sorted(os.listdir(record_dir)): + zf.write(os.path.join(record_dir, name), arcname=name) + + with pytest.raises(NotImplementedError, match="Multi-segment"): + wfdb.rdrecord(archive_path) + with pytest.raises(NotImplementedError, match="Multi-segment"): + wfdb.rdheader(archive_path, rd_segments=True) + + +def test_read_archive_unsupported_locations(): + with pytest.raises(ValueError, match="pn_dir"): + wfdb.rdrecord("100.wfdb", pn_dir="mitdb") + with pytest.raises(NotImplementedError, match="cloud storage"): + wfdb.rdrecord("s3://bucket/100.wfdb") + + +def test_create_archive_compression(tmp_path): + """Files are stored uncompressed. + + This way, part of a signal file can be read without decompressing all of it, and + reading an archive doesn't require a decompression library. + """ + archive_path = archive_sample_record("sample-data/100", tmp_path) + + with zipfile.ZipFile(archive_path) as zf: + compress_types = {info.filename: info.compress_type for info in zf.infolist()} + + assert compress_types == { + "100.hea": zipfile.ZIP_STORED, + "100.dat": zipfile.ZIP_STORED, + } + + +def test_create_archive_replaces_existing(tmp_path): + archive_path = os.fspath(tmp_path / "100.wfdb") + create_archive( + record_files("sample-data/100") + ["sample-data/100.atr"], archive_path + ) + + create_archive(record_files("sample-data/100"), archive_path) + + with zipfile.ZipFile(archive_path) as zf: + assert sorted(zf.namelist()) == ["100.dat", "100.hea"] + + +def test_archive_write_error_keeps_existing_archive(tmp_path): + """An error while writing leaves the existing archive unchanged. + + No temporary files are left behind either. + """ + archive_path = archive_sample_record("sample-data/100", tmp_path) + with open(archive_path, "rb") as f: + original = f.read() + + with pytest.raises(RuntimeError): + with WFDBArchive(archive_path, mode="w") as archive: + archive.add_file("sample-data/test01_00s.hea") + raise RuntimeError("failed while writing") + + with open(archive_path, "rb") as f: + assert f.read() == original + assert os.listdir(tmp_path) == ["100.wfdb"] + + +def test_archive_mode_checks(tmp_path): + archive_path = archive_sample_record("sample-data/100", tmp_path) + + with WFDBArchive(archive_path, mode="r") as archive: + with pytest.raises(ValueError, match="cannot be written to"): + archive.write("extra.txt", b"data") + + with WFDBArchive(os.fspath(tmp_path / "new.wfdb"), mode="w") as archive: + archive.write("new.hea", b"new 0\n") + with pytest.raises(ValueError, match="already written"): + archive.write("new.hea", b"new 0\n") + with pytest.raises(ValueError, match="cannot be read from"): + archive.open("new.hea") + + with pytest.raises(ValueError, match="mode"): + WFDBArchive(archive_path, mode="a") + + +def test_archive_header_named_after_archive(tmp_path): + """The record's header must be named after the archive. + + The archive is not saved otherwise. + """ + with pytest.raises(ValueError, match="must be named after the archive"): + wfdb.wrsamp("rec", wfdb_archive=tmp_path / "other.wfdb", **write_kwargs()) + with pytest.raises(ValueError, match="must be named after the archive"): + create_archive(record_files("sample-data/100"), tmp_path / "other.wfdb") + + assert os.listdir(tmp_path) == [] + + +@pytest.mark.parametrize( + "file_name", ["sub/100.dat", "../100.dat", ".100.dat", "100 a.dat", "é.dat", ""] +) +def test_archive_file_names(tmp_path, file_name): + """Files are at the top level, with ASCII names that follow the WFDB rules.""" + with WFDBArchive(tmp_path / "100.wfdb", mode="w") as archive: + with pytest.raises(ValueError, match="at the top level"): + archive.write(file_name, b"x") + archive.write("100_1.dat", b"x") + + +def test_archive_holds_one_record(tmp_path): + """Only one header file can be written to an archive. + + This way, it is clear which record the archive holds. Other files can be added. + """ + archive_path = os.fspath(tmp_path / "100.wfdb") + + with pytest.raises(ValueError, match="holds a single record"): + create_archive( + record_files("sample-data/100") + ["sample-data/test01_00s.hea"], + archive_path, + ) + assert os.listdir(tmp_path) == [] + + with WFDBArchive(archive_path, mode="w") as archive: + for path in record_files("sample-data/100"): + archive.add_file(path) + archive.add_file("sample-data/100.atr") + with pytest.raises(ValueError, match="already contains the header"): + archive.write("other.hea", b"other 0\n") + + with zipfile.ZipFile(archive_path) as zf: + assert sorted(zf.namelist()) == ["100.atr", "100.dat", "100.hea"] + + +def write_kwargs(fmt="16"): + rng = np.random.default_rng(1234) + return dict( + fs=250, + units=["mV", "mV"], + sig_name=["I", "II"], + p_signal=rng.standard_normal((1000, 2)), + fmt=[fmt, fmt], + ) + + +@pytest.mark.parametrize("fmt", ["16", "24", "32", "516"]) +def test_wrsamp_archive_matches_files(fmt, tmp_path): + """Writing to an archive gives the same result as writing separate files.""" + kwargs = write_kwargs(fmt) + archive_path = os.fspath(tmp_path / "archive" / "rec.wfdb") + files_dir = tmp_path / "files" + files_dir.mkdir() + + wfdb.wrsamp("rec", wfdb_archive=archive_path, **kwargs) + wfdb.wrsamp("rec", write_dir=os.fspath(files_dir), **kwargs) + + for physical in [True, False]: + record = wfdb.rdrecord(archive_path, physical=physical) + expected = wfdb.rdrecord(os.fspath(files_dir / "rec"), physical=physical) + assert_records_equal(record, expected) + + +def test_wrsamp_archive_object(tmp_path): + """An open WFDBArchive is written to and left open, with no separate files.""" + archive_path = os.fspath(tmp_path / "rec.wfdb") + + with WFDBArchive(archive_path, mode="w") as archive: + wfdb.wrsamp( + "rec", + write_dir=os.fspath(tmp_path), + wfdb_archive=archive, + **write_kwargs(), + ) + archive.write("notes.txt", b"extra file") + + assert os.listdir(tmp_path) == ["rec.wfdb"] + with zipfile.ZipFile(archive_path) as zf: + assert sorted(zf.namelist()) == ["notes.txt", "rec.dat", "rec.hea"] + assert wfdb.rdrecord(archive_path).sig_len == 1000 + + +def test_record_wrsamp_archive_round_trip(tmp_path): + """A record read from an archive can be written to a new archive.""" + source = archive_sample_record("sample-data/100", tmp_path) + record = wfdb.rdrecord(source, physical=False) + copy_path = os.fspath(tmp_path / "copy" / "100.wfdb") + + record.wrsamp(wfdb_archive=copy_path) + + assert_records_equal( + wfdb.rdrecord(copy_path, physical=False), + wfdb.rdrecord("sample-data/100", physical=False), + ) + + +def test_wrsamp_archive_replaces_existing(tmp_path): + archive_path = archive_sample_record("sample-data/100", tmp_path) + + wfdb.wrsamp("100", wfdb_archive=archive_path, **write_kwargs()) + + with zipfile.ZipFile(archive_path) as zf: + assert sorted(zf.namelist()) == ["100.dat", "100.hea"] + assert wfdb.rdrecord(archive_path).sig_name == ["I", "II"] + + +def test_wrsamp_archive_error_writes_nothing(tmp_path, monkeypatch): + """A failed write leaves no archive or temporary file. + + This is the case even if the failure happens after the header was written. + """ + + def fail(*args, **kwargs): + raise RuntimeError("failed while writing the dat file") + + monkeypatch.setattr(wfdb.io._signal, "wr_dat_file", fail) + + with pytest.raises(RuntimeError, match="dat file"): + wfdb.wrsamp( + "rec", + wfdb_archive=os.fspath(tmp_path / "rec.wfdb"), + **write_kwargs(), + ) + + assert os.listdir(tmp_path) == [] + + +def test_wrsamp_two_records_in_one_archive(tmp_path): + """Writing a second record into an archive fails. + + The archive is not saved. + """ + archive_path = os.fspath(tmp_path / "rec.wfdb") + + with pytest.raises(ValueError, match="holds a single record"): + with WFDBArchive(archive_path, mode="w") as archive: + wfdb.wrsamp("rec", wfdb_archive=archive, **write_kwargs()) + wfdb.wrsamp("other", wfdb_archive=archive, **write_kwargs()) + + assert os.listdir(tmp_path) == [] + + +def test_wrsamp_archive_invalid_argument(tmp_path): + archive_path = archive_sample_record("sample-data/100", tmp_path) + + with pytest.raises(TypeError, match="wfdb_archive"): + wfdb.wrsamp("rec", wfdb_archive=123, **write_kwargs()) + + with WFDBArchive(archive_path, mode="r") as archive: + with pytest.raises(ValueError, match="cannot be written to"): + wfdb.wrsamp("rec", wfdb_archive=archive, **write_kwargs()) + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX file permissions") +@pytest.mark.parametrize("umask", [0o022, 0o027]) +def test_archive_permissions(umask, tmp_path, monkeypatch): + """A new archive gets the same permissions as other new files. + + A replaced archive keeps its permissions. The umask is never changed, since that + would affect files created by other threads at the same time. + """ + archive_path = tmp_path / "rec.wfdb" + other_path = tmp_path / "other.txt" + old_umask = os.umask(umask) + try: + other_path.write_text("") + + def fail(mask): + raise AssertionError("the umask was changed") + + monkeypatch.setattr(os, "umask", fail) + wfdb.wrsamp("rec", wfdb_archive=os.fspath(archive_path), **write_kwargs()) + finally: + monkeypatch.undo() + os.umask(old_umask) + assert archive_path.stat().st_mode == other_path.stat().st_mode + + archive_path.chmod(0o640) + wfdb.wrsamp("rec", wfdb_archive=os.fspath(archive_path), **write_kwargs()) + assert archive_path.stat().st_mode & 0o777 == 0o640 + + +def test_archive_temporary_file(tmp_path): + """While an archive is written, its temporary file is named after it. + + If the archive is discarded without being closed, the temporary file is removed and + nothing is saved. + """ + archive = WFDBArchive(os.fspath(tmp_path / "rec.wfdb"), mode="w") + archive.write("rec.hea", b"rec 0\n") + (tmp_file,) = os.listdir(tmp_path) + assert tmp_file.startswith(".rec.wfdb.") and tmp_file.endswith(".tmp") + + del archive + gc.collect() + + assert os.listdir(tmp_path) == [] + + +def test_rdrecord_opens_archive_once(tmp_path, monkeypatch): + """The header and the signals are read from the same open archive.""" + archive_path = archive_sample_record("sample-data/100", tmp_path) + opened = [] + original_init = WFDBArchive.__init__ + + def counting_init(self, *args, **kwargs): + opened.append(args) + original_init(self, *args, **kwargs) + + monkeypatch.setattr(WFDBArchive, "__init__", counting_init) + + wfdb.rdrecord(archive_path, sampto=100) + + assert len(opened) == 1 + + +def test_rdann_archive(tmp_path): + """Annotations read from an archive match the annotation file on disk.""" + archive_path = os.fspath(tmp_path / "100.wfdb") + create_archive( + record_files("sample-data/100") + ["sample-data/100.atr"], + archive_path, + ) + + ann = wfdb.rdann(archive_path, "atr", sampfrom=1000, sampto=50000) + expected = wfdb.rdann("sample-data/100", "atr", sampfrom=1000, sampto=50000) + + assert ann.record_name == "100" + assert ann.extension == "atr" + assert ann.fs == expected.fs + np.testing.assert_array_equal(ann.sample, expected.sample) + assert ann.symbol == expected.symbol + assert ann.aux_note == expected.aux_note + + +def test_rdann_archive_header_named_differently(tmp_path): + """Annotation files are named after the record's header, not the archive.""" + archive_path = os.fspath(tmp_path / "renamed.wfdb") + zip_files(record_files("sample-data/100") + ["sample-data/100.atr"], archive_path) + + ann = wfdb.rdann(archive_path, "atr") + + assert ann.record_name == "100" + assert len(ann.sample) == len(wfdb.rdann("sample-data/100", "atr").sample) + + +def test_rdann_archive_errors(tmp_path): + archive_path = archive_sample_record("sample-data/100", tmp_path) + + with pytest.raises(FileNotFoundError, match="100.atr"): + wfdb.rdann(archive_path, "atr") + with pytest.raises(ValueError, match="pn_dir"): + wfdb.rdann(archive_path, "atr", pn_dir="mitdb") + + +def test_create_archive_old_file_dates(tmp_path): + """Files dated before 1980, which ZIP can't represent, can be archived.""" + for path in record_files("sample-data/100"): + copy = tmp_path / os.path.basename(path) + shutil.copyfile(path, copy) + os.utime(copy, (0, 0)) + archive_path = os.fspath(tmp_path / "archive" / "100.wfdb") + + create_archive([tmp_path / "100.hea", tmp_path / "100.dat"], archive_path) + + expected = wfdb.rdrecord("sample-data/100", sampto=100) + record = wfdb.rdrecord(archive_path, sampto=100) + np.testing.assert_array_equal(record.p_signal, expected.p_signal) + + +def test_archive_abort_removes_temporary_file(tmp_path): + """The temporary file is removed even if closing the zip file fails.""" + archive = WFDBArchive(os.fspath(tmp_path / "rec.wfdb"), mode="w") + archive.write("rec.hea", b"rec 0\n") + + zip_file = archive._zipfile + real_close = zip_file.close + + def fail(): + # Close the file, as when writing the end of the zip file fails + zip_file.close = real_close + real_close() + raise OSError("disk full") + + zip_file.close = fail + + with pytest.raises(OSError, match="disk full"): + archive.abort() + + assert os.listdir(tmp_path) == [] + + +def test_closed_archive(tmp_path): + """Using a closed archive raises a clear error.""" + archive_path = archive_sample_record("sample-data/100", tmp_path) + + archive = WFDBArchive(archive_path, mode="r") + archive.close() + for call in [ + archive.namelist, + lambda: archive.exists("100.hea"), + lambda: archive.file_size("100.hea"), + lambda: archive.open("100.hea"), + archive.read_header, + ]: + with pytest.raises(ValueError, match="is closed"): + call() + + archive = WFDBArchive(os.fspath(tmp_path / "new.wfdb"), mode="w") + archive.close() + with pytest.raises(ValueError, match="is closed"): + archive.write("new.hea", b"new 0\n") + with pytest.raises(ValueError, match="is closed"): + archive.add_file("sample-data/100.hea") + with pytest.raises(ValueError, match="is closed"): + wfdb.wrsamp("rec", wfdb_archive=archive, **write_kwargs()) + + +def test_wrsamp_archive_write_dir(tmp_path, monkeypatch): + """A relative archive path is relative to `write_dir`; an absolute one isn't.""" + write_dir = tmp_path / "write_dir" + write_dir.mkdir() + cwd = tmp_path / "cwd" + cwd.mkdir() + monkeypatch.chdir(cwd) + + wfdb.wrsamp( + "rec", write_dir=os.fspath(write_dir), wfdb_archive="rec.wfdb", **write_kwargs() + ) + record = wfdb.rdrecord(os.fspath(write_dir / "rec.wfdb"), physical=False) + record.wrsamp( + write_dir=os.fspath(write_dir), wfdb_archive=os.path.join("copy", "rec.wfdb") + ) + record.wrsamp( + write_dir=os.fspath(write_dir), + wfdb_archive=os.fspath(tmp_path / "abs" / "rec.wfdb"), + ) + + assert os.listdir(cwd) == [] + assert sorted(os.listdir(write_dir)) == ["copy", "rec.wfdb"] + assert os.listdir(write_dir / "copy") == ["rec.wfdb"] + assert os.listdir(tmp_path / "abs") == ["rec.wfdb"] + + +def test_archive_extension_case(tmp_path): + """The .wfdb extension is recognized in any case.""" + archive_path = os.fspath(tmp_path / "100.WFDB") + create_archive(record_files("sample-data/100"), archive_path) + + assert os.listdir(tmp_path) == ["100.WFDB"] + record = wfdb.rdrecord(archive_path, sampto=100) + assert record.record_name == "100" + assert wfdb.rdheader(archive_path).sig_name == ["MLII", "V5"] + with WFDBArchive(archive_path) as archive: + assert archive.record_name == "100" + assert archive.path == archive_path + + +def test_archive_write_bytes_like(tmp_path, monkeypatch): + """`write` streams any bytes-like object, with the same entries as `writestr`. + + ZIP64 is used for large files: the limit is lowered so that a small file needs it. + """ + signal = np.arange(1000, dtype=np.int16) + archive_path = os.fspath(tmp_path / "rec.wfdb") + with WFDBArchive(archive_path, mode="w") as archive: + archive.write("rec.hea", b"rec 1 250 1000\n") + archive.write("array.dat", signal) + archive.write("view.dat", memoryview(signal.tobytes())) + archive.write("bytearray.dat", bytearray(b"abc")) + monkeypatch.setattr(zipfile, "ZIP64_LIMIT", 100) + archive.write("zip64.dat", signal) + monkeypatch.undo() + + with zipfile.ZipFile(archive_path) as zf: + assert zf.read("array.dat") == signal.tobytes() + assert zf.read("view.dat") == signal.tobytes() + assert zf.read("bytearray.dat") == b"abc" + assert zf.read("zip64.dat") == signal.tobytes() + for info in zf.infolist(): + assert info.date_time[0] >= 2020 + assert info.external_attr == 0o600 << 16 + assert zf.getinfo("rec.hea").compress_type == zipfile.ZIP_STORED + assert zf.getinfo("array.dat").compress_type == zipfile.ZIP_STORED