From dcce27cd8e8cce4960b34e37e68c53c7ed41bb95 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Tue, 27 Oct 2026 16:24:49 -0400 Subject: [PATCH 01/16] Add support for reading .wfdb archive files. A .wfdb archive is a ZIP file holding the header, signal and annotation files of one record. rdrecord, rdsamp, rdheader and rdann now read from an archive when given its path, including the .wfdb extension. The header is the file named after the archive, or the only header in it. Signal and annotation files are read by the names that follow from the header, and a missing file is an error. rdrecord reads the header and signals from a single open archive. Multi-segment records and archives in cloud storage or on PhysioNet are not supported yet. An archive holds a single record: WFDBArchive refuses to write a second header file, so an archive can't be written that rdrecord cannot read. Annotation and other files can be added alongside the record. --- src/wfdb/io/_signal.py | 66 ++++++- src/wfdb/io/annotation.py | 46 ++++- src/wfdb/io/archive.py | 334 ++++++++++++++++++++++++++++++++ src/wfdb/io/record.py | 143 ++++++++++++-- tests/test_archive.py | 393 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 953 insertions(+), 29 deletions(-) create mode 100644 src/wfdb/io/archive.py create mode 100644 tests/test_archive.py diff --git a/src/wfdb/io/_signal.py b/src/wfdb/io/_signal.py index 33cfc517..f9debafc 100644 --- a/src/wfdb/io/_signal.py +++ b/src/wfdb/io/_signal.py @@ -1099,6 +1099,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). @@ -1146,6 +1147,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 ------- @@ -1243,6 +1247,7 @@ def _rd_segment( sampto=sampto, no_file=no_file, sig_data=sig_data, + wfdb_archive=wfdb_archive, ) # Copy over the wanted signals @@ -1267,6 +1272,7 @@ def _rd_dat_signals( sampto, no_file=False, sig_data=None, + wfdb_archive=None, ): """Read all signals from a WFDB dat file. @@ -1304,6 +1310,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 ------- @@ -1365,10 +1374,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: @@ -1606,7 +1622,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 @@ -1631,6 +1649,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 ------- @@ -1662,8 +1683,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( @@ -1811,6 +1840,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. @@ -1837,6 +1867,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 ------- @@ -1861,10 +1894,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") @@ -2489,7 +2526,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. @@ -2509,6 +2552,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 ------- @@ -2522,9 +2568,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 2321fb5e..6c3e3b0f 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: @@ -1893,7 +1894,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'. @@ -1935,8 +1939,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): @@ -1952,7 +1961,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( @@ -1992,7 +2005,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, @@ -2101,6 +2114,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 path.endswith(ARCHIVE_EXTENSION): + 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: + directory = os.path.dirname(os.path.abspath(path)) + os.makedirs(directory, exist_ok=True) + fd, self._tmp_path = tempfile.mkstemp( + dir=directory, + prefix=f".{os.path.basename(path)}.", + suffix=".tmp", + ) + os.close(fd) + try: + self._zipfile = zipfile.ZipFile(self._tmp_path, mode="w") + 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.""" + return self._zipfile.namelist() + + def exists(self, file_name): + """Check whether a file is in the archive.""" + return file_name in self._zipfile.namelist() + + 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 = [n for n in self.namelist() if n.endswith(".hea")] + 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() + if self.file_size(file_name) > MAX_HEADER_SIZE: + raise ValueError( + f"Header file '{file_name}' in archive '{self.path}' is " + f"larger than {MAX_HEADER_SIZE} bytes" + ) + with self.open(file_name) as f: + # Don't trust the size stored in the archive + 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. + + Parameters + ---------- + file_name : str + The name of the file in the archive. + data : bytes + The content of the file. + + """ + self._check_new_file(file_name) + self._zipfile.writestr(file_name, data, compress_type=_compress_type(file_name)) + + 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, compress_type=_compress_type(file_name) + ) + + def close(self): + """Close the archive. + + In write mode, this moves the new archive into place. + """ + if self._zipfile is None: + return + self._finalizer.detach() + try: + self._zipfile.close() + if self._tmp_path is not None: + os.chmod(self._tmp_path, _new_file_mode(self.path)) + 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): + 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 self.exists(file_name): + raise ValueError(f"'{file_name}' was already written to the archive") + if file_name.endswith(".hea"): + headers = [n for n in self.namelist() if n.endswith(".hea")] + 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 _require_mode(self, mode): + 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.""" + if zip_file is not None: + zip_file.close() + if tmp_path is not None and os.path.exists(tmp_path): + os.remove(tmp_path) + + +def _new_file_mode(path): + """Return the permissions for a new file at `path`. + + These are the permissions of the existing file, or the default for new files (based + on the umask). + """ + try: + return stat.S_IMODE(os.stat(path).st_mode) + except FileNotFoundError: + # The umask can only be read by setting it + umask = os.umask(0) + os.umask(umask) + return 0o666 & ~umask + + +def _compress_type(file_name): + """Return the compression type for a file in an archive. + + Header files are compressed, but signal files are stored as they are, so that + reading part of a signal doesn't require decompressing all of it. + """ + if file_name.endswith(".hea"): + return zipfile.ZIP_DEFLATED + return zipfile.ZIP_STORED + + +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. 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.""" + return os.fspath(record_name).endswith(ARCHIVE_EXTENSION) diff --git a/src/wfdb/io/record.py b/src/wfdb/io/record.py index 107b0fc9..aaab837d 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, is_archive_path # -------------- WFDB Signal Calibration and Classification ---------- # @@ -1759,7 +1761,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 +1786,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 +1828,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 +1910,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 +1958,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 +2035,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 +2105,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 +2138,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 +2215,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 diff --git a/tests/test_archive.py b/tests/test_archive.py new file mode 100644 index 00000000..aa435cdc --- /dev/null +++ b/tests/test_archive.py @@ -0,0 +1,393 @@ +import gc +import os +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 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"] + + create_archive(record_files("sample-data/test01_00s"), archive_path) + + record = wfdb.rdrecord(archive_path) + expected = wfdb.rdrecord("sample-data/test01_00s") + assert_records_equal(record, expected) + + +def test_read_archive_header_named_differently(tmp_path): + """A single header can be read even if it isn't named after the archive.""" + archive_path = os.fspath(tmp_path / "renamed.wfdb") + create_archive(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) + + +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): + """Headers are compressed, and signal files are stored uncompressed. + + This way, part of a signal file can be read without decompressing all of it. + """ + 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_DEFLATED, + "100.dat": zipfile.ZIP_STORED, + } + + +def test_create_archive_replaces_existing(tmp_path): + archive_path = os.fspath(tmp_path / "record.wfdb") + create_archive(record_files("sample-data/test01_00s"), 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_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 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") + create_archive( + 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") From b9b6f46de729e2c2a6b16be03aae23743e866c76 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Tue, 27 Oct 2026 16:25:04 -0400 Subject: [PATCH 02/16] Add support for writing .wfdb archive files. wrsamp and Record.wrsamp take a wfdb_archive argument: a path creates a new archive containing the header and dat files of the record, and an open WFDBArchive lets several files be written to one archive. A new archive is written to a temporary file named after it, and moved into place only when writing succeeds, so a failed write leaves no partial archive and an existing archive at the same path is either fully replaced or left unchanged. The new archive gets the permissions of the archive it replaces, or the default permissions for new files. --- src/wfdb/io/_header.py | 20 ++++-- src/wfdb/io/_signal.py | 35 ++++++++--- src/wfdb/io/archive.py | 27 ++++++++ src/wfdb/io/record.py | 38 +++++++++--- tests/test_archive.py | 136 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 235 insertions(+), 21 deletions(-) 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 f9debafc..96730773 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"): @@ -2309,6 +2320,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. @@ -2332,6 +2344,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 ------- @@ -2458,8 +2473,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, @@ -2468,7 +2484,9 @@ def wr_dat_file( ) with sf: sf.write(d_signal) - return + if wfdb_archive is not None: + wfdb_archive.write(file_name, target.getvalue()) + return else: raise ValueError( @@ -2488,8 +2506,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, b_write.tobytes()) + else: + with open(file_path, "wb") as f: + b_write.tofile(f) def describe_list_indices(full_list): diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index c6bcba56..eca3770e 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -13,6 +13,7 @@ import tempfile import weakref import zipfile +from contextlib import contextmanager ARCHIVE_EXTENSION = ".wfdb" @@ -332,3 +333,29 @@ def create_archive(file_paths, archive_path): def is_archive_path(record_name): """Check whether a record name refers to a .wfdb archive.""" return os.fspath(record_name).endswith(ARCHIVE_EXTENSION) + + +@contextmanager +def archive_for_writing(wfdb_archive): + """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. + + """ + 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)): + with WFDBArchive(wfdb_archive, 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 aaab837d..79071804 100644 --- a/src/wfdb/io/record.py +++ b/src/wfdb/io/record.py @@ -11,7 +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, is_archive_path +from wfdb.io.archive import WFDBArchive, archive_for_writing, is_archive_path # -------------- WFDB Signal Calibration and Classification ---------- # @@ -872,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 @@ -882,6 +882,11 @@ 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 WFDBArchive must be opened in + write mode, and is left open. Returns ------- @@ -898,13 +903,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) 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. @@ -2937,6 +2945,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. @@ -2992,6 +3001,11 @@ 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 WFDBArchive must be opened in + write mode, and is left open. Returns ------- @@ -3016,6 +3030,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 @@ -3120,7 +3138,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 index aa435cdc..b59592d8 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -317,6 +317,142 @@ def test_archive_holds_one_record(tmp_path): 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") +def test_archive_permissions(tmp_path): + """A new archive gets the same permissions as other new files. + + A replaced archive keeps its permissions. + """ + archive_path = tmp_path / "rec.wfdb" + other_path = tmp_path / "other.txt" + other_path.write_text("") + + wfdb.wrsamp("rec", wfdb_archive=os.fspath(archive_path), **write_kwargs()) + 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. From b62c548a8a0be69b8057a9622ee7e7fbf6ce8e81 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Tue, 27 Oct 2026 16:25:04 -0400 Subject: [PATCH 03/16] Document .wfdb archive support. --- docs/changes.rst | 3 +++ docs/io.rst | 46 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/docs/changes.rst b/docs/changes.rst index 45e5401f..4f978541 100644 --- a/docs/changes.rst +++ b/docs/changes.rst @@ -74,6 +74,9 @@ Unreleased **Allow None for annotations without notes in wrann** ``wrann`` now accepts ``None`` in ``aux_note`` for annotations without a note, as the docstring says, and writes it the same way as an empty string. It used to raise a ``TypeError``. +**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..ec78ef6f 100644 --- a/docs/io.rst +++ b/docs/io.rst @@ -34,6 +34,52 @@ 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. +Header files are compressed, while signal files are stored without +compression so that part of a signal can be read without decompressing +the whole file. + +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 -------------------- From 4d284fa7b73cd3144ecbc64bbfa4174b75398cc3 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 21:56:55 -0400 Subject: [PATCH 04/16] Allow archiving files dated before 1980. ZIP can't represent dates before 1980, and ZipFile.write raises a ValueError for such files by default. This happens with files from reproducible builds (SOURCE_DATE_EPOCH=0), for example. These files now get the earliest date ZIP supports. --- src/wfdb/io/archive.py | 6 +++++- tests/test_archive.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index eca3770e..b148f120 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -85,7 +85,11 @@ def __init__(self, path, mode="r"): ) os.close(fd) try: - self._zipfile = zipfile.ZipFile(self._tmp_path, mode="w") + # 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 diff --git a/tests/test_archive.py b/tests/test_archive.py index b59592d8..b470927a 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -1,5 +1,6 @@ import gc import os +import shutil import zipfile import numpy as np @@ -527,3 +528,18 @@ def test_rdann_archive_errors(tmp_path): 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) From 68ffa914d3332e7efde30c5c16c977dbe7d24ee2 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 21:57:25 -0400 Subject: [PATCH 05/16] Don't change the umask when writing an archive. To give a new archive the default permissions for new files, the umask was read by setting it to 0 and back. The umask is process-wide, so files that another thread created in between came out world-writable. The temporary file is now created with os.open and mode 0o666, so the operating system applies the umask, instead of with tempfile.mkstemp, which uses 0o600. The permissions are only copied when an existing archive is replaced. --- src/wfdb/io/archive.py | 49 ++++++++++++++++++++++++------------------ tests/test_archive.py | 21 +++++++++++++----- 2 files changed, 44 insertions(+), 26 deletions(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index b148f120..ce08a7e6 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -9,8 +9,8 @@ """ import os +import secrets import stat -import tempfile import weakref import zipfile from contextlib import contextmanager @@ -76,14 +76,7 @@ def __init__(self, path, mode="r"): except zipfile.BadZipFile as e: raise ValueError(f"Not a valid WFDB archive: {path}") from e else: - directory = os.path.dirname(os.path.abspath(path)) - os.makedirs(directory, exist_ok=True) - fd, self._tmp_path = tempfile.mkstemp( - dir=directory, - prefix=f".{os.path.basename(path)}.", - suffix=".tmp", - ) - os.close(fd) + 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 @@ -227,7 +220,14 @@ def close(self): try: self._zipfile.close() if self._tmp_path is not None: - os.chmod(self._tmp_path, _new_file_mode(self.path)) + # 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) @@ -286,19 +286,26 @@ def _discard(zip_file, tmp_path): os.remove(tmp_path) -def _new_file_mode(path): - """Return the permissions for a new file at `path`. +def _create_temp_file(path): + """Create an empty temporary file next to `path`, and return its path. - These are the permissions of the existing file, or the default for new files (based - on the umask). + 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. """ - try: - return stat.S_IMODE(os.stat(path).st_mode) - except FileNotFoundError: - # The umask can only be read by setting it - umask = os.umask(0) - os.umask(umask) - return 0o666 & ~umask + 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 _compress_type(file_name): diff --git a/tests/test_archive.py b/tests/test_archive.py index b470927a..f53ffc3d 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -437,16 +437,27 @@ def test_wrsamp_archive_invalid_argument(tmp_path): @pytest.mark.skipif(os.name == "nt", reason="POSIX file permissions") -def test_archive_permissions(tmp_path): +@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. + 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" - other_path.write_text("") - - wfdb.wrsamp("rec", wfdb_archive=os.fspath(archive_path), **write_kwargs()) + 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) From 8bc3c452b0977fca138ea4480e5efb2ff2f41322 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 21:58:29 -0400 Subject: [PATCH 06/16] Remove the temporary file even if closing the archive fails. When an archive is discarded, the temporary file was only removed if closing the zip file succeeded, so a failure (such as a full disk) left it behind. --- src/wfdb/io/archive.py | 17 ++++++++++++----- tests/test_archive.py | 22 ++++++++++++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index ce08a7e6..0728fd2a 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -8,6 +8,7 @@ decompressing the whole file. """ +import contextlib import os import secrets import stat @@ -279,11 +280,17 @@ def _require_mode(self, mode): def _discard(zip_file, tmp_path): - """Close a zip file and remove its temporary file, if there is one.""" - if zip_file is not None: - zip_file.close() - if tmp_path is not None and os.path.exists(tmp_path): - os.remove(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): diff --git a/tests/test_archive.py b/tests/test_archive.py index f53ffc3d..300a2866 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -554,3 +554,25 @@ def test_create_archive_old_file_dates(tmp_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) == [] From bd6f2382087b70002a452a9d78c24492fc0b3b03 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 21:58:47 -0400 Subject: [PATCH 07/16] Raise a clear error when a closed archive is used. Using a WFDBArchive after closing it failed with an AttributeError on None. It now raises a ValueError saying that the archive is closed. --- src/wfdb/io/archive.py | 9 ++++++++- tests/test_archive.py | 26 ++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 0728fd2a..cc928910 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -110,11 +110,12 @@ def record_name(self): 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.""" - return file_name in self._zipfile.namelist() + return file_name in self.namelist() def file_size(self, file_name): """Return the uncompressed size of a file in the archive, in bytes.""" @@ -249,6 +250,7 @@ def abort(self): self._tmp_path = None def _getinfo(self, file_name): + self._require_open() try: return self._zipfile.getinfo(file_name) except KeyError: @@ -270,7 +272,12 @@ def _check_new_file(self, file_name): f"header '{headers[0]}'" ) + 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( diff --git a/tests/test_archive.py b/tests/test_archive.py index 300a2866..67b6b49d 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -576,3 +576,29 @@ def fail(): 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()) From 36b7d6a415cd7b1558158fe5d89b36e63154cc8d Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 21:59:16 -0400 Subject: [PATCH 08/16] Write archives given as a relative path into write_dir. When wfdb_archive was a path, write_dir was silently ignored, so the archive ended up in the current directory. A relative archive path is now relative to write_dir, like the files of a record written without an archive. An absolute path is used as it is. --- src/wfdb/io/archive.py | 8 ++++++-- src/wfdb/io/record.py | 12 +++++++----- tests/test_archive.py | 22 ++++++++++++++++++++++ 3 files changed, 35 insertions(+), 7 deletions(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index cc928910..0377f448 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -361,7 +361,7 @@ def is_archive_path(record_name): @contextmanager -def archive_for_writing(wfdb_archive): +def archive_for_writing(wfdb_archive, write_dir=""): """Get a WFDBArchive to write a record to. Parameters @@ -370,6 +370,9 @@ def archive_for_writing(wfdb_archive): 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: @@ -378,7 +381,8 @@ def archive_for_writing(wfdb_archive): wfdb_archive._require_mode("w") yield wfdb_archive elif isinstance(wfdb_archive, (str, os.PathLike)): - with WFDBArchive(wfdb_archive, mode="w") as archive: + path = os.path.join(write_dir, wfdb_archive) + with WFDBArchive(path, mode="w") as archive: yield archive else: raise TypeError( diff --git a/src/wfdb/io/record.py b/src/wfdb/io/record.py index 79071804..52a879b3 100644 --- a/src/wfdb/io/record.py +++ b/src/wfdb/io/record.py @@ -885,8 +885,9 @@ def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None): 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 WFDBArchive must be opened in - write mode, and is left open. + 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 ------- @@ -903,7 +904,7 @@ def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None): checksums[ch] = old_val self.checksum = checksums - with archive_for_writing(wfdb_archive) as archive: + 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) @@ -3004,8 +3005,9 @@ def wrsamp( 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 WFDBArchive must be opened in - write mode, and is left open. + 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 ------- diff --git a/tests/test_archive.py b/tests/test_archive.py index 67b6b49d..0f43887f 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -602,3 +602,25 @@ def test_closed_archive(tmp_path): 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="copy.wfdb") + record.wrsamp( + write_dir=os.fspath(write_dir), wfdb_archive=os.fspath(tmp_path / "abs.wfdb") + ) + + assert os.listdir(cwd) == [] + assert sorted(os.listdir(write_dir)) == ["copy.wfdb", "rec.wfdb"] + assert os.path.exists(tmp_path / "abs.wfdb") From 848aa41b85fad4b374747e41105ac6c4c00a27c6 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 21:59:38 -0400 Subject: [PATCH 09/16] Recognize the .wfdb extension in any case. A path such as '100.WFDB' was treated as a plain record name, and WFDBArchive added another '.wfdb' to it. The extension is now compared without regard to case. --- src/wfdb/io/archive.py | 11 +++++++---- tests/test_archive.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 0377f448..8ec3c5d7 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -45,7 +45,7 @@ class WFDBArchive: ---------- path : str or os.PathLike The path of the archive. The .wfdb extension is added if it is - missing. + missing (the extension is not case-sensitive). mode : str, optional 'r' to read an existing archive (default), or 'w' to write a new one. @@ -63,7 +63,7 @@ def __init__(self, path, mode="r"): raise ValueError("mode must be 'r' or 'w'") path = os.fspath(path) - if not path.endswith(ARCHIVE_EXTENSION): + if not is_archive_path(path): path += ARCHIVE_EXTENSION self.path = path self.mode = mode @@ -356,8 +356,11 @@ def create_archive(file_paths, archive_path): def is_archive_path(record_name): - """Check whether a record name refers to a .wfdb archive.""" - return os.fspath(record_name).endswith(ARCHIVE_EXTENSION) + """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 diff --git a/tests/test_archive.py b/tests/test_archive.py index 0f43887f..ba762a7e 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -624,3 +624,17 @@ def test_wrsamp_archive_write_dir(tmp_path, monkeypatch): assert os.listdir(cwd) == [] assert sorted(os.listdir(write_dir)) == ["copy.wfdb", "rec.wfdb"] assert os.path.exists(tmp_path / "abs.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 From 435e21a57958599453d3bded9c050e166ef3ebf7 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 22:06:52 -0400 Subject: [PATCH 10/16] Avoid copying signal files in memory when writing to an archive. Signal files were copied before being added to an archive: tobytes() for dat files, and BytesIO.getvalue() for FLAC files. WFDBArchive.write now streams any bytes-like object into the archive, so the array (or the FLAC buffer) is passed without a copy. The archive entries are written as before, and ZIP64 is used for files of 2 GiB or more. --- src/wfdb/io/_signal.py | 4 ++-- src/wfdb/io/archive.py | 20 +++++++++++++++++--- tests/test_archive.py | 28 ++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/src/wfdb/io/_signal.py b/src/wfdb/io/_signal.py index 96730773..e7a34a4c 100644 --- a/src/wfdb/io/_signal.py +++ b/src/wfdb/io/_signal.py @@ -2485,7 +2485,7 @@ def wr_dat_file( with sf: sf.write(d_signal) if wfdb_archive is not None: - wfdb_archive.write(file_name, target.getvalue()) + wfdb_archive.write(file_name, target.getbuffer()) return else: @@ -2507,7 +2507,7 @@ def wr_dat_file( # Write the bytes to the file if wfdb_archive is not None: - wfdb_archive.write(file_name, b_write.tobytes()) + wfdb_archive.write(file_name, np.ascontiguousarray(b_write)) else: with open(file_path, "wb") as f: b_write.tofile(f) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 8ec3c5d7..7bbe1ef1 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -12,6 +12,7 @@ import os import secrets import stat +import time import weakref import zipfile from contextlib import contextmanager @@ -185,16 +186,29 @@ def read_header(self): 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 - The content of the file. + 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) - self._zipfile.writestr(file_name, data, compress_type=_compress_type(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.compress_type = _compress_type(file_name) + 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. diff --git a/tests/test_archive.py b/tests/test_archive.py index ba762a7e..8c2a385e 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -638,3 +638,31 @@ def test_archive_extension_case(tmp_path): 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_DEFLATED + assert zf.getinfo("array.dat").compress_type == zipfile.ZIP_STORED From b943a663aeae3225876541e98ef682b98805da42 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Wed, 28 Oct 2026 22:10:27 -0400 Subject: [PATCH 11/16] Remove duplicate checks in WFDBArchive. - read_header no longer checks the declared size of the header before reading it: the bounded read already enforces the limit, whatever size the archive declares. - The scan for header files is in one helper, used by header_file_name and _check_new_file. - exists() looks the file up with getinfo, instead of building the list of all files on every call. --- src/wfdb/io/archive.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 7bbe1ef1..337becc8 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -116,7 +116,12 @@ def namelist(self): def exists(self, file_name): """Check whether a file is in the archive.""" - return file_name in self.namelist() + 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.""" @@ -143,7 +148,7 @@ def header_file_name(self): """ file_name = self.record_name + ".hea" if not self.exists(file_name): - headers = [n for n in self.namelist() if n.endswith(".hea")] + headers = self._header_names() if len(headers) != 1: raise FileNotFoundError( f"Could not find '{file_name}' in archive '{self.path}'" @@ -168,13 +173,9 @@ def read_header(self): """ file_name = self.header_file_name() - if self.file_size(file_name) > MAX_HEADER_SIZE: - raise ValueError( - f"Header file '{file_name}' in archive '{self.path}' is " - f"larger than {MAX_HEADER_SIZE} bytes" - ) with self.open(file_name) as f: - # Don't trust the size stored in the archive + # 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( @@ -278,7 +279,7 @@ def _check_new_file(self, file_name): if self.exists(file_name): raise ValueError(f"'{file_name}' was already written to the archive") if file_name.endswith(".hea"): - headers = [n for n in self.namelist() if n.endswith(".hea")] + headers = self._header_names() if headers: raise ValueError( f"Cannot add '{file_name}' to '{self.path}': an archive " @@ -286,6 +287,10 @@ def _check_new_file(self, file_name): f"header '{headers[0]}'" ) + def _header_names(self): + """Return the names of the header files in the archive.""" + return [n for n in self.namelist() if n.endswith(".hea")] + def _require_open(self): if self._zipfile is None: raise ValueError(f"Archive '{self.path}' is closed") From ecaa8b5b2c449616356d9fa6a50788a6795f9e53 Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Thu, 29 Oct 2026 11:25:20 -0400 Subject: [PATCH 12/16] Update src/wfdb/io/archive.py Co-authored-by: Luke Payyapilli --- src/wfdb/io/archive.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 337becc8..89563625 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -289,7 +289,9 @@ def _check_new_file(self, file_name): def _header_names(self): """Return the names of the header files in the archive.""" - return [n for n in self.namelist() if n.endswith(".hea")] + # Only files at the top level belong to the record. macOS, for example, + # adds '__MACOSX/._100.hea' when it zips '100.hea'. + return [n for n in self.namelist() if n.endswith(".hea") and "/" not in n] def _require_open(self): if self._zipfile is None: From 605bfb655193939e80da41632c830024d1bad0ac Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Thu, 29 Oct 2026 11:25:29 -0400 Subject: [PATCH 13/16] Update tests/test_archive.py Co-authored-by: Luke Payyapilli --- tests/test_archive.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_archive.py b/tests/test_archive.py index 8c2a385e..93e33c54 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -163,6 +163,24 @@ def test_read_archive_several_headers(tmp_path): wfdb.rdrecord(archive_path) +def test_read_archive_ignores_files_in_folders(tmp_path): + """Only header files at the top level of an archive are considered. + + macOS, for example, adds '__MACOSX/._100.hea' when it zips '100.hea', which + shouldn't count as a second header. + """ + 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(f"__MACOSX/._{file_name}", b"\x00\x05\x16\x07") + + record = wfdb.rdrecord(archive_path, sampto=100) + + assert record.record_name == "100" + + def test_read_archive_missing_dat_file(tmp_path): """A dat file named in the header but missing from the archive is an error. From 5db9e681e3f65ea1cf3ae816455d9ea36d57244f Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Thu, 29 Oct 2026 11:26:55 -0400 Subject: [PATCH 14/16] Ignore macOS metadata files when looking for headers in an archive. --- src/wfdb/io/archive.py | 17 +++++++++++++---- tests/test_archive.py | 12 +++++++----- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 89563625..9997e535 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -10,6 +10,7 @@ import contextlib import os +import re import secrets import stat import time @@ -23,6 +24,10 @@ # small, so anything bigger is treated as a damaged or malicious file. MAX_HEADER_SIZE = 1024 * 1024 +# A header file at the top level of an archive, named after a record (record names +# are made of letters, digits, hyphens and underscores) +_HEADER_NAME = re.compile(r"[-\w]+\.hea") + class WFDBArchive: """A .wfdb archive opened for reading or writing. @@ -288,10 +293,14 @@ def _check_new_file(self, file_name): ) def _header_names(self): - """Return the names of the header files in the archive.""" - # Only files at the top level belong to the record. macOS, for example, - # adds '__MACOSX/._100.hea' when it zips '100.hea'. - return [n for n in self.namelist() if n.endswith(".hea") and "/" not in n] + """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: diff --git a/tests/test_archive.py b/tests/test_archive.py index 93e33c54..b40f0752 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -163,18 +163,20 @@ def test_read_archive_several_headers(tmp_path): wfdb.rdrecord(archive_path) -def test_read_archive_ignores_files_in_folders(tmp_path): - """Only header files at the top level of an archive are considered. +@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. - macOS, for example, adds '__MACOSX/._100.hea' when it zips '100.hea', which - shouldn't count as a second header. + 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(f"__MACOSX/._{file_name}", b"\x00\x05\x16\x07") + zf.writestr(prefix + file_name, b"\x00\x05\x16\x07") record = wfdb.rdrecord(archive_path, sampto=100) From 1fc8868c7b7fde78b6e2936ea2c8c1ba9fb3c6fc Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Thu, 29 Oct 2026 11:27:37 -0400 Subject: [PATCH 15/16] Store all files in an archive uncompressed. --- docs/io.rst | 6 +++--- src/wfdb/io/archive.py | 23 +++++------------------ tests/test_archive.py | 26 ++++++++++++++++++++++---- 3 files changed, 30 insertions(+), 25 deletions(-) diff --git a/docs/io.rst b/docs/io.rst index ec78ef6f..e5671ad6 100644 --- a/docs/io.rst +++ b/docs/io.rst @@ -40,9 +40,9 @@ 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. -Header files are compressed, while signal files are stored without -compression so that part of a signal can be read without decompressing -the whole file. +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 diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 9997e535..2587e2c0 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -3,9 +3,10 @@ A .wfdb archive is a ZIP file containing the files of a single WFDB record: its header file, the signal files listed in the header, and optionally annotation files. -Header files are compressed, while signal files are stored without -compression so that a range of samples can be read without -decompressing the whole file. +Files are stored without compression, so that a range of samples can +be read without decompressing the whole signal file, and so that +reading an archive doesn't require a decompression library. Archives +made by other tools may compress their files, and can still be read. """ import contextlib @@ -208,7 +209,6 @@ def write(self, file_name, data): data = memoryview(data).cast("B") # The same file information that ZipFile.writestr uses info = zipfile.ZipInfo(file_name, date_time=time.localtime()[:6]) - info.compress_type = _compress_type(file_name) 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( @@ -227,9 +227,7 @@ def add_file(self, path): """ file_name = os.path.basename(path) self._check_new_file(file_name) - self._zipfile.write( - path, arcname=file_name, compress_type=_compress_type(file_name) - ) + self._zipfile.write(path, arcname=file_name) def close(self): """Close the archive. @@ -352,17 +350,6 @@ def _create_temp_file(path): raise FileExistsError(f"Could not create a temporary file next to '{path}'") -def _compress_type(file_name): - """Return the compression type for a file in an archive. - - Header files are compressed, but signal files are stored as they are, so that - reading part of a signal doesn't require decompressing all of it. - """ - if file_name.endswith(".hea"): - return zipfile.ZIP_DEFLATED - return zipfile.ZIP_STORED - - def create_archive(file_paths, archive_path): """Create a .wfdb archive from the files of a record on disk. diff --git a/tests/test_archive.py b/tests/test_archive.py index b40f0752..bf632472 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -183,6 +183,23 @@ def test_read_archive_ignores_macos_metadata(tmp_path, prefix): 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. @@ -251,9 +268,10 @@ def test_read_archive_unsupported_locations(): def test_create_archive_compression(tmp_path): - """Headers are compressed, and signal files are stored uncompressed. + """Files are stored uncompressed. - This way, part of a signal file can be read without decompressing all of it. + 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) @@ -261,7 +279,7 @@ def test_create_archive_compression(tmp_path): compress_types = {info.filename: info.compress_type for info in zf.infolist()} assert compress_types == { - "100.hea": zipfile.ZIP_DEFLATED, + "100.hea": zipfile.ZIP_STORED, "100.dat": zipfile.ZIP_STORED, } @@ -684,5 +702,5 @@ def test_archive_write_bytes_like(tmp_path, monkeypatch): 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_DEFLATED + assert zf.getinfo("rec.hea").compress_type == zipfile.ZIP_STORED assert zf.getinfo("array.dat").compress_type == zipfile.ZIP_STORED From c7dba93564e6bc8c408bf621c7fe194c438a3d5f Mon Sep 17 00:00:00 2001 From: Tom Pollard Date: Thu, 29 Oct 2026 11:41:24 -0400 Subject: [PATCH 16/16] Only write archive files with WFDB names, named after the archive. --- docs/io.rst | 4 ++- src/wfdb/io/archive.py | 45 +++++++++++++++++++++++----- tests/test_archive.py | 68 ++++++++++++++++++++++++++++++++---------- 3 files changed, 92 insertions(+), 25 deletions(-) diff --git a/docs/io.rst b/docs/io.rst index e5671ad6..fedfb539 100644 --- a/docs/io.rst +++ b/docs/io.rst @@ -39,7 +39,9 @@ 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. +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. diff --git a/src/wfdb/io/archive.py b/src/wfdb/io/archive.py index 2587e2c0..3c775cd0 100644 --- a/src/wfdb/io/archive.py +++ b/src/wfdb/io/archive.py @@ -25,9 +25,12 @@ # small, so anything bigger is treated as a damaged or malicious file. MAX_HEADER_SIZE = 1024 * 1024 -# A header file at the top level of an archive, named after a record (record names -# are made of letters, digits, hyphens and underscores) -_HEADER_NAME = re.compile(r"[-\w]+\.hea") +# The name of a file at the top level of an archive: an ASCII WFDB name made of +# letters, digits, hyphens and underscores, with an extension +_FILE_NAME = re.compile(r"[-\w]+(\.[-\w]+)*", re.ASCII) + +# The name of a header file at the top level of an archive +_HEADER_NAME = re.compile(r"[-\w]+\.hea", re.ASCII) class WFDBArchive: @@ -41,8 +44,11 @@ class WFDBArchive: the default permissions for new files. An archive holds a single record, so only one header file can be - written to it. Annotation files and other files can be added - alongside the record's header and signal files. + written to it, and it must be named after the archive (for example + 'record.hea' in 'record.wfdb'). Annotation files and other files can + be added alongside the record's header and signal files. All files + are at the top level of the archive, and their names are made of + ASCII letters, digits, hyphens and underscores, with an extension. Use the object as a context manager so that it is always closed. An archive in write mode that is discarded without being closed is not @@ -236,6 +242,12 @@ def close(self): """ 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() @@ -279,9 +291,15 @@ def _getinfo(self, file_name): 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 file_name.endswith(".hea"): + if _HEADER_NAME.fullmatch(file_name): headers = self._header_names() if headers: raise ValueError( @@ -290,6 +308,17 @@ def _check_new_file(self, file_name): 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. @@ -359,8 +388,8 @@ def create_archive(file_paths, archive_path): 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. An existing archive at this - path is replaced. + The path of the archive to create, which must be named after the + record. An existing archive at this path is replaced. Examples -------- diff --git a/tests/test_archive.py b/tests/test_archive.py index bf632472..d15b485b 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -21,6 +21,13 @@ def record_files(record_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")) @@ -131,17 +138,19 @@ def test_read_replaced_archive(tmp_path): archive_path = archive_sample_record("sample-data/100", tmp_path) assert wfdb.rdrecord(archive_path).sig_name == ["MLII", "V5"] - create_archive(record_files("sample-data/test01_00s"), archive_path) + wfdb.wrsamp("100", wfdb_archive=archive_path, **write_kwargs()) - record = wfdb.rdrecord(archive_path) - expected = wfdb.rdrecord("sample-data/test01_00s") - assert_records_equal(record, expected) + 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.""" + """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") - create_archive(record_files("sample-data/100"), archive_path) + zip_files(record_files("sample-data/100"), archive_path) record = wfdb.rdrecord(archive_path, sampto=100) @@ -285,8 +294,10 @@ def test_create_archive_compression(tmp_path): def test_create_archive_replaces_existing(tmp_path): - archive_path = os.fspath(tmp_path / "record.wfdb") - create_archive(record_files("sample-data/test01_00s"), archive_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) @@ -331,6 +342,30 @@ def test_archive_mode_checks(tmp_path): 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. @@ -559,10 +594,7 @@ def test_rdann_archive(tmp_path): 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") - create_archive( - record_files("sample-data/100") + ["sample-data/100.atr"], - archive_path, - ) + zip_files(record_files("sample-data/100") + ["sample-data/100.atr"], archive_path) ann = wfdb.rdann(archive_path, "atr") @@ -654,14 +686,18 @@ def test_wrsamp_archive_write_dir(tmp_path, monkeypatch): "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="copy.wfdb") record.wrsamp( - write_dir=os.fspath(write_dir), wfdb_archive=os.fspath(tmp_path / "abs.wfdb") + 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.wfdb", "rec.wfdb"] - assert os.path.exists(tmp_path / "abs.wfdb") + 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):