Fix file close issue.

Author: tompollard

tests/test_archive.py

@@ -54,11 +54,11 @@ def temp_record():
         dat_path = os.path.join(tmpdir, record_basename + ".dat")
         archive_path = os.path.join(tmpdir, record_basename + ".wfdb")
 
-        WFDBArchive.create_archive(
-            None,
-            file_list=[hea_path, dat_path],
-            output_path=archive_path,
-        )
+        with WFDBArchive(record_name=record_basename, mode="w") as archive:
+            archive.create_archive(
+                file_list=[hea_path, dat_path],
+                output_path=archive_path,
+            )
 
         yield {
             "record_name": os.path.join(tmpdir, record_basename),

wfdb/io/archive.py

@@ -1,3 +1,4 @@
+import io
 import os
 import shutil
 import zipfile
@@ -44,10 +45,9 @@ def __init__(self, record_name, mode="r"):
             self.zipfile = zipfile.ZipFile(self.archive_path, mode="r")
 
         elif mode == "w":
-            # Initialize an empty archive on disk
+            # Create archive file if needed
             if not os.path.exists(self.archive_path):
-                with zipfile.ZipFile(self.archive_path, mode="w"):
-                    pass  # Just create the file
+                WFDBArchive.make_archive_file([], self.archive_path)
             self.zipfile = zipfile.ZipFile(self.archive_path, mode="a")
 
     def __enter__(self):
@@ -62,6 +62,15 @@ def exists(self, filename):
         """
         return self.zipfile and filename in self.zipfile.namelist()
 
+    @staticmethod
+    def make_archive_file(file_list, output_path):
+        with zipfile.ZipFile(output_path, mode="w") as zf:
+            for file in file_list:
+                compress = zipfile.ZIP_DEFLATED
+                zf.write(
+                    file, arcname=os.path.basename(file), compress_type=compress
+                )
+
     @contextmanager
     def open(self, filename, mode="r"):
         """
@@ -73,8 +82,6 @@ def open(self, filename, mode="r"):
                 if "b" in mode:
                     yield f
                 else:
-                    import io
-
                     yield io.TextIOWrapper(f)
         else:
             raise FileNotFoundError(
@@ -97,7 +104,7 @@ def write(self, filename, data):
             self.zipfile.writestr(filename, data)
             return
 
-        # If already opened in read or append mode, use the replace-then-move trick
+        # If already opened in read or append mode, use replace-then-move
         tmp_path = self.archive_path + ".tmp"
         with zipfile.ZipFile(self.archive_path, mode="r") as zin:
             with zipfile.ZipFile(tmp_path, mode="w") as zout:
@@ -114,16 +121,13 @@ def create_archive(self, file_list, output_path=None):
         If output_path is not specified, uses self.archive_path.
         """
         output_path = output_path or self.archive_path
-        with zipfile.ZipFile(output_path, mode="w") as zf:
-            for file in file_list:
-                compress = (
-                    zipfile.ZIP_STORED
-                    if file.endswith((".hea", ".hea.json", ".hea.yml"))
-                    else zipfile.ZIP_DEFLATED
-                )
-                zf.write(
-                    file, arcname=os.path.basename(file), compress_type=compress
-                )
+        WFDBArchive.make_archive_file(file_list, output_path)
+
+        # If this archive object points to the archive, reload the zipfile in append mode
+        if output_path == self.archive_path:
+            if self.zipfile:
+                self.zipfile.close()
+            self.zipfile = zipfile.ZipFile(self.archive_path, mode="a")
 
 
 def get_archive(record_base_name, mode="r"):