fix: close zst tar resources safely

Author: DylanWelzel

biothings/utils/common.py

@@ -19,15 +19,17 @@
 import os.path
 import pickle
 import random
+import shutil
 import string
 import sys
 import tarfile
+import tempfile
 import time
 import types
 import urllib.parse
 import warnings
 from collections import UserDict, UserList
-from contextlib import contextmanager
+from contextlib import closing, contextmanager
 from datetime import date, datetime, timezone
 from functools import partial
 from itertools import islice
@@ -188,22 +190,35 @@ def anyfile(infile, mode="r"):
     # tarfile handling. works for zst in Python >= 3.14
     if lower_version_zst or tarfile.is_tarfile(infile):
         if lower_version_zst:
-            f = open(infile, "rb")
-            dctx = zstd.ZstdDecompressor()
-            reader = dctx.stream_reader(f)
-            tar_file = tarfile.open(fileobj=reader, mode="r|")  # streaming mode
-        else:
-            tar_file = tarfile.open(infile, mode)
-
-        extracted = None
+            with open(infile, "rb") as compressed_file:
+                dctx = zstd.ZstdDecompressor()
+                with closing(dctx.stream_reader(compressed_file)) as reader:
+                    with tarfile.open(fileobj=reader, mode="r|") as tar_file:
+                        for member in tar_file:
+                            if member.name == rawfile:
+                                extracted = tar_file.extractfile(member)
+                                break
+                        else:
+                            extracted = None
+
+                        # Keep the returned file readable after closing the tar and zst streams.
+                        if extracted is not None:
+                            with extracted:
+                                spooled_file = tempfile.SpooledTemporaryFile(  # pylint: disable=consider-using-with
+                                    max_size=1024 * 1024
+                                )
+                                shutil.copyfileobj(extracted, spooled_file)
+                            spooled_file.seek(0)
+
+            # extracted member is not a regular file or link
+            if extracted is None:
+                raise Exception("invalid target file: must be a regular file or a link")
+
+            return spooled_file
+
+        tar_file = tarfile.open(infile, mode)  # pylint: disable=consider-using-with
         try:
-            if lower_version_zst:
-                for member in tar_file:
-                    if member.name == rawfile:
-                        extracted = tar_file.extractfile(member)
-                        break
-            else:
-                extracted = tar_file.extractfile(rawfile)
+            extracted = tar_file.extractfile(rawfile)
         except KeyError:
             # provided rawfile does not appear in the tarball
             tar_file.close()
@@ -214,9 +229,6 @@ def anyfile(infile, mode="r"):
             tar_file.close()
             raise Exception("invalid target file: must be a regular file or a link")
 
-        if lower_version_zst:
-            return extracted
-
         return io.TextIOWrapper(extracted)
 
     if filetype == ".gz":