Skip to content

Commit 8971d79

Browse files
committed
Support writing directly with wrsamp.
1 parent 388dfb0 commit 8971d79

File tree

4 files changed

+67
-14
lines changed

4 files changed

+67
-14
lines changed

tests/test_archive.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ def test_wfdb_archive_inline_round_trip():
9898
sig = (np.random.randn(sig_len, 2) * 1000).astype(np.float32)
9999

100100
# Create archive inline using context manager
101-
with WFDBArchive(record_basename, mode="w") as wfdb_archive:
101+
with WFDBArchive(record_path, mode="w") as wfdb_archive:
102102
wrsamp(
103103
record_name=record_basename,
104104
fs=fs,

wfdb/io/_signal.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1710,6 +1710,16 @@ def _rd_dat_file(
17101710

17111711
# Local file or .wfdb archive
17121712
if wfdb_archive is not None:
1713+
# If the exact file name isn't found, look for any .dat file
1714+
if not wfdb_archive.exists(file_name):
1715+
dat_files = [f for f in wfdb_archive.zipfile.namelist()
1716+
if f.endswith('.dat')]
1717+
if not dat_files:
1718+
raise FileNotFoundError(
1719+
f"No dat file found in archive for {file_name}"
1720+
)
1721+
file_name = dat_files[0] # Use the first dat file found
1722+
17131723
with wfdb_archive.open(file_name, "rb") as fp:
17141724
fp.seek(start_byte)
17151725
sig_data = util.fromfile(

wfdb/io/archive.py

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,11 @@ def __init__(self, record_name, mode="r"):
3131
'r' for read (default), 'w' for write.
3232
"""
3333
self.record_name = record_name
34-
self.archive_path = f"{record_name}.wfdb"
34+
# Only append .wfdb if it's not already there
35+
if record_name.endswith('.wfdb'):
36+
self.archive_path = record_name
37+
else:
38+
self.archive_path = f"{record_name}.wfdb"
3539
self.zipfile = None
3640
self.mode = mode
3741

@@ -41,14 +45,24 @@ def __init__(self, record_name, mode="r"):
4145
f"Archive not found: {self.archive_path}"
4246
)
4347
if not zipfile.is_zipfile(self.archive_path):
44-
raise ValueError(f"Invalid WFDB archive: {self.archive_path}")
45-
self.zipfile = zipfile.ZipFile(self.archive_path, mode="r")
48+
raise ValueError(
49+
f"Invalid WFDB archive: {self.archive_path}"
50+
)
51+
self.zipfile = zipfile.ZipFile(
52+
self.archive_path, mode="r"
53+
)
4654

4755
elif mode == "w":
4856
# Create archive file if needed
4957
if not os.path.exists(self.archive_path):
50-
WFDBArchive.make_archive_file([], self.archive_path)
51-
self.zipfile = zipfile.ZipFile(self.archive_path, mode="a")
58+
# Create the directory if it doesn't exist
59+
os.makedirs(os.path.dirname(os.path.abspath(self.archive_path)), exist_ok=True)
60+
WFDBArchive.make_archive_file(
61+
[], self.archive_path
62+
)
63+
self.zipfile = zipfile.ZipFile(
64+
self.archive_path, mode="a"
65+
)
5266

5367
def __enter__(self):
5468
return self
@@ -85,7 +99,8 @@ def open(self, filename, mode="r"):
8599
yield io.TextIOWrapper(f)
86100
else:
87101
raise FileNotFoundError(
88-
f"Could not find '{filename}' as loose file or inside '{self.archive_path}'."
102+
f"Could not find '{filename}' as loose file or inside "
103+
f"'{self.archive_path}'."
89104
)
90105

91106
def close(self):

wfdb/io/record.py

Lines changed: 35 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
from wfdb.io import _header, _signal, _url, download, header, util
1212
from wfdb.io._coreio import CLOUD_PROTOCOLS
13-
from wfdb.io.archive import get_archive
13+
from wfdb.io.archive import get_archive, WFDBArchive
1414

1515
# -------------- WFDB Signal Calibration and Classification ---------- #
1616

@@ -911,6 +911,10 @@ def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None):
911911
of the uniform signal (d_signal).
912912
write_dir : str, optional
913913
The directory in which to write the files.
914+
wfdb_archive : str or WFDBArchive, optional
915+
If provided, writes the record to a .wfdb archive. Can be either:
916+
- A string path to the archive file (e.g., 'record.wfdb')
917+
- A WFDBArchive object for more advanced usage
914918
915919
Returns
916920
-------
@@ -927,6 +931,15 @@ def wrsamp(self, expanded=False, write_dir="", wfdb_archive=None):
927931
checksums[ch] = old_val
928932
self.checksum = checksums
929933

934+
# Handle wfdb_archive parameter
935+
if wfdb_archive:
936+
if isinstance(wfdb_archive, str):
937+
# If a string path is provided, create a WFDBArchive object
938+
from wfdb.io.archive import get_archive
939+
wfdb_archive = get_archive(wfdb_archive, mode="w")
940+
elif not isinstance(wfdb_archive, WFDBArchive):
941+
raise TypeError("wfdb_archive must be either a string path or WFDBArchive object")
942+
930943
# Perform field validity and cohesion checks, and write the
931944
# header file.
932945
self.wrheader(
@@ -1975,7 +1988,6 @@ def rdrecord(
19751988
Option used to stream data from Physionet. The Physionet
19761989
database directory from which to find the required record files.
19771990
eg. For record '100' in 'http://physionet.org/content/mitdb'
1978-
pn_dir='mitdb'.
19791991
m2s : bool, optional
19801992
Used when reading multi-segment records. Specifies whether to
19811993
directly return a WFDB MultiRecord object (False), or to convert
@@ -2033,15 +2045,19 @@ def rdrecord(
20332045
--------
20342046
>>> record = wfdb.rdrecord('sample-data/test01_00s', sampfrom=800,
20352047
channels=[1, 3])
2036-
20372048
"""
20382049
wfdb_archive = None
20392050
is_wfdb_archive = record_name.endswith(".wfdb")
20402051

20412052
if is_wfdb_archive:
20422053
record_base = record_name[:-5] # remove ".wfdb"
20432054
wfdb_archive = get_archive(record_base)
2044-
hea_file = os.path.basename(record_base) + ".hea"
2055+
2056+
# Find any .hea file in the archive
2057+
hea_files = [f for f in wfdb_archive.zipfile.namelist() if f.endswith('.hea')]
2058+
if not hea_files:
2059+
raise FileNotFoundError(f"No header file found in archive {record_name}")
2060+
hea_file = hea_files[0] # Use the first header file found
20452061

20462062
import tempfile
20472063

@@ -2954,6 +2970,10 @@ def wrsamp(
29542970
setting both `base_date` and `base_time`.
29552971
write_dir : str, optional
29562972
The directory in which to write the files.
2973+
wfdb_archive : str or WFDBArchive, optional
2974+
If provided, writes the record to a .wfdb archive. Can be either:
2975+
- A string path to the archive file (e.g., 'record.wfdb')
2976+
- A WFDBArchive object for more advanced usage
29572977
29582978
Returns
29592979
-------
@@ -2978,6 +2998,10 @@ def wrsamp(
29782998
>>> # Write a local WFDB record (manually inserting fields)
29792999
>>> wfdb.wrsamp('ecgrecord', fs = 250, units=['mV', 'mV'],
29803000
sig_name=['I', 'II'], p_signal=signals, fmt=['16', '16'])
3001+
>>> # Write to a .wfdb archive using a string path
3002+
>>> wfdb.wrsamp('ecgrecord', fs = 250, units=['mV', 'mV'],
3003+
sig_name=['I', 'II'], p_signal=signals, fmt=['16', '16'],
3004+
wfdb_archive='ecgrecord.wfdb')
29813005
29823006
"""
29833007
# Check for valid record name
@@ -3082,10 +3106,14 @@ def wrsamp(
30823106
else:
30833107
expanded = False
30843108

3109+
# Handle wfdb_archive parameter
30853110
if wfdb_archive:
3086-
wfdb_archive = get_archive(
3087-
os.path.join(write_dir, record_name), mode="w"
3088-
)
3111+
if isinstance(wfdb_archive, str):
3112+
# If a string path is provided, create a WFDBArchive object
3113+
from wfdb.io.archive import get_archive
3114+
wfdb_archive = get_archive(wfdb_archive, mode="w")
3115+
elif not isinstance(wfdb_archive, WFDBArchive):
3116+
raise TypeError("wfdb_archive must be either a string path or WFDBArchive object")
30893117

30903118
# Write the record files - header and associated dat
30913119
record.wrsamp(

0 commit comments

Comments
 (0)