xixu-me/tzst
1
0 Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Improve robustness in TzstArchive handling

Browse Source 
Enhanced error handling and validation in TzstArchive methods, including better checks for file object initialization and compressed stream creation. Updated compression level validation to use dynamic max level from zstd library. Simplified and clarified comments and error messages for unsupported modes and extraction operations.
This commit is contained in:
xixu-me committed 2025-06-09 10:41:09 +08:00
1 parent 36ffd4fd2b
commit 320ed0a27a
1 file changed
+91 -103
Download Patch File Download Diff File Expand all files Collapse all files
src/tzst/core.py
+91 -103
View File
@@ -105,7 +105,7 @@ def _handle_file_conflict(
# Convert string resolution to enum if needed
if isinstance(resolution, str):
try:
resolution = ConflictResolution(resolution)
resolution = ConflictResolution(resolution.lower())
except ValueError:
# Invalid string, fallback to ASK for interactive handling
resolution = ConflictResolution.ASK
@@ -130,7 +130,6 @@ def _handle_file_conflict(
elif resolution == ConflictResolution.EXIT:
return resolution, None
else:
# Unknown resolution, default to REPLACE for robustness
return ConflictResolution.REPLACE, target_path
@@ -176,19 +175,20 @@ class TzstArchive:
)
# Validate compression level
if not 1 <= compression_level <= 22:
if not 1 <= compression_level <= zstd.ZstdCompressor.max_level():
raise ValueError(
f"Invalid compression level '{compression_level}'. Must be between 1 and 22."
f"Invalid compression level '{compression_level}'. "
f"Must be between 1 and {zstd.ZstdCompressor.max_level()}."
)
# Check for unsupported modes immediately - provide clear documentation
# Check for unsupported modes immediately
if mode.startswith("a"):
raise NotImplementedError(
"Append mode is not currently supported for .tzst/.tar.zst archives. "
"This would require decompressing the entire archive, adding new files, "
"and recompressing, which is complex and potentially slow for large archives. "
"Alternatives: 1) Create multiple archives, 2) Recreate the archive with all files, "
"3) Use standard tar format for append operations, then compress separately."
"This involves complex operations (decompress, add, recompress) "
"and is slow for large archives. Consider alternatives like creating "
"multiple archives, recreating the archive, or using standard tar "
"for appends and compressing separately."
)
def __enter__(self):
@@ -212,24 +212,28 @@ class TzstArchive:
"""
try:
if self.mode.startswith("r"):
# Read mode
self._fileobj = open(self.filename, "rb")
dctx = zstd.ZstdDecompressor()
if self.streaming:
# Streaming mode - use stream reader directly (memory efficient)
# Note: This may limit some tarfile operations that require seeking
self._compressed_stream = dctx.stream_reader(self._fileobj)
if not isinstance(self._compressed_stream, io.IOBase):
raise TzstArchiveError(
"Failed to create a valid compressed stream for reading."
)
self._tarfile = tarfile.open(
fileobj=self._compressed_stream, mode="r|"
fileobj=self._compressed_stream,
mode="r|", # type: ignore[arg-type]
)
else:
# Buffer mode - decompress to memory buffer for random access
# Better compatibility but higher memory usage for large archives
if self._fileobj is None:
raise TzstArchiveError(
"File object not initialized for reading."
)
decompressed_chunks = []
with dctx.stream_reader(self._fileobj) as reader:
while True:
chunk = reader.read(8192)
chunk = reader.read(io.DEFAULT_BUFFER_SIZE)
if not chunk:
break
decompressed_chunks.append(chunk)
@@ -240,25 +244,20 @@ class TzstArchive:
)
elif self.mode.startswith("w"):
# Write mode - use streaming compression
self._fileobj = open(self.filename, "wb")
cctx = zstd.ZstdCompressor(
level=self.compression_level, write_content_size=True
)
if self._fileobj is None:
raise TzstArchiveError("File object not initialized for writing.")
self._compressed_stream = cctx.stream_writer(self._fileobj)
self._tarfile = tarfile.open(fileobj=self._compressed_stream, mode="w|")
elif self.mode.startswith("a"):
# Append mode - for tar.zst, this is complex as we need to decompress,
# add files, and recompress. For simplicity, we'll raise an error for now.
raise NotImplementedError(
"Append mode is not currently supported for .tzst/.tar.zst archives. "
"This would require decompressing the entire archive, adding new files, "
"and recompressing, which is complex and potentially slow for large archives. "
"Alternatives: 1) Create multiple archives, 2) Recreate the archive with all files, "
"3) Use standard tar format for append operations, then compress separately."
)
else:
raise ValueError(f"Invalid mode: {self.mode}")
if not isinstance(self._compressed_stream, io.IOBase):
raise TzstArchiveError(
"Failed to create a valid compressed stream for writing."
)
self._tarfile = tarfile.open(fileobj=self._compressed_stream, mode="w|") # type: ignore[arg-type]
# 'a' mode already handled by raising NotImplementedError earlier
# else case for invalid mode also handled earlier
except Exception as e:
self.close()
if "zstd" in str(e).lower():
@@ -340,7 +339,7 @@ class TzstArchive:
- 'data': Safe filter for cross-platform data archives (default, recommended)
- 'tar': Honor most tar features but block dangerous ones
- 'fully_trusted': Honor all metadata (use only for trusted archives)
- None: Use default behavior (may show deprecation warning in Python 3.12+)
- None: Use default behavior (Python 3.12+ may warn about this)
- callable: Custom filter function
Warning:
@@ -368,24 +367,34 @@ class TzstArchive:
# Specific member extraction not supported in streaming mode
raise RuntimeError(
"Extracting specific members is not supported in streaming mode. "
"Please use non-streaming mode for selective extraction, or extract all files."
) # Prepare extraction arguments - different parameters for extract vs extractall
"Use non-streaming mode or extract all files."
)
try:
if member:
# extract() accepts set_attrs, numeric_owner, and filter
extract_kwargs = {
"set_attrs": set_attrs,
"numeric_owner": numeric_owner,
"filter": filter,
}
self._tarfile.extract(member, path=extract_path, **extract_kwargs)
# Ensure correct types for extract_kwargs
_filter_arg: (
str
| Callable[[tarfile.TarInfo, str], tarfile.TarInfo | None]
| None
) = filter
self._tarfile.extract( # type: ignore[call-arg]
member,
path=extract_path,
set_attrs=set_attrs, # type: ignore[call-arg]
numeric_owner=numeric_owner, # type: ignore[call-arg]
filter=_filter_arg, # type: ignore[call-arg]
)
else:
# extractall() only accepts numeric_owner and filter (no set_attrs)
extractall_kwargs = {
"numeric_owner": numeric_owner,
"filter": filter,
}
self._tarfile.extractall(path=extract_path, **extractall_kwargs)
_filter_arg_all: (
str
| Callable[[tarfile.TarInfo, str], tarfile.TarInfo | None]
| None
) = filter
self._tarfile.extractall( # type: ignore[call-arg]
path=extract_path,
numeric_owner=numeric_owner, # type: ignore[call-arg]
filter=_filter_arg_all, # type: ignore[call-arg]
)
except (tarfile.StreamError, OSError) as e:
if self.streaming and (
"seeking" in str(e).lower() or "stream" in str(e).lower()
@@ -421,7 +430,9 @@ class TzstArchive:
members: list[tarfile.TarInfo] | None = None,
*,
numeric_owner: bool = False,
filter: str | Callable | None = "data",
filter: str
| Callable[[tarfile.TarInfo, str], tarfile.TarInfo | None]
| None = "data",
):
"""
Extract all members from the archive.
@@ -468,15 +479,17 @@ class TzstArchive:
extract_path.mkdir(parents=True, exist_ok=True)
try:
# extractall() accepts numeric_owner, filter, and members parameters
extractall_kwargs = {
"numeric_owner": numeric_owner,
"filter": filter,
}
if members is not None:
extractall_kwargs["members"] = members
_filter_arg: (
str | Callable[[tarfile.TarInfo, str], tarfile.TarInfo | None] | None
) = filter
_members_arg: list[tarfile.TarInfo] | None = members
self._tarfile.extractall(path=extract_path, **extractall_kwargs)
self._tarfile.extractall(
path=extract_path,
members=_members_arg,
numeric_owner=numeric_owner,
filter=_filter_arg, # type: ignore[call-arg]
)
except (tarfile.StreamError, OSError) as e:
if self.streaming and (
"seeking" in str(e).lower() or "stream" in str(e).lower()
@@ -585,7 +598,7 @@ class TzstArchive:
if fileobj:
# Read the entire file to verify decompression
while True:
chunk = fileobj.read(8192)
chunk = fileobj.read(io.DEFAULT_BUFFER_SIZE)
if not chunk:
break
return True
@@ -603,62 +616,37 @@ def create_archive(
use_temp_file: bool = True,
) -> None:
"""
Create a new .tzst archive with atomic file operations.
Create a .tzst/.tar.zst archive.
Args:
archive_path: Path for the new archive
files: List of files/directories to add
archive_path: Path to the archive file to be created
files: List of files or directories to add to the archive
compression_level: Zstandard compression level (1-22)
use_temp_file: If True, create archive in temporary file first, then move
to final location for atomic operation
See Also:
:meth:`TzstArchive.add`: Method for adding files to an open archive
use_temp_file: If True, create archive in a temporary file first,
then move for atomicity (recommended).
"""
# Validate compression level
if not 1 <= compression_level <= 22:
raise ValueError(
f"Invalid compression level '{compression_level}'. Must be between 1 and 22."
)
archive_path_obj = Path(archive_path)
archive_path = Path(archive_path)
# Ensure archive has correct extension
if archive_path.suffix.lower() not in [".tzst", ".zst"]:
if archive_path.suffix.lower() == ".tar":
archive_path = archive_path.with_suffix(".tar.zst")
else:
archive_path = archive_path.with_suffix(archive_path.suffix + ".tzst")
# Use temporary file for atomic operation if requested
if use_temp_file:
temp_fd = None
temp_path = None
temp_dir = archive_path_obj.parent
temp_dir.mkdir(parents=True, exist_ok=True)
fd, temp_archive_name = tempfile.mkstemp(
suffix=".tzst.tmp", prefix=archive_path_obj.name + "_", dir=str(temp_dir)
)
os.close(fd)
temp_archive_path = Path(temp_archive_name)
try:
# Create temporary file in same directory as target for atomic move
temp_fd, temp_path_str = tempfile.mkstemp(
suffix=".tmp", prefix=f".{archive_path.name}.", dir=archive_path.parent
)
os.close(temp_fd) # Close file descriptor, we'll open with TzstArchive
temp_path = Path(temp_path_str)
# Create archive in temporary location
_create_archive_impl(temp_path, files, compression_level)
# Atomic move to final location
temp_path.replace(archive_path)
_create_archive_impl(temp_archive_path, files, compression_level)
archive_path_obj.parent.mkdir(parents=True, exist_ok=True)
temp_archive_path.rename(archive_path_obj)
except Exception:
# Clean up temporary file on error
if temp_path and temp_path.exists():
try:
temp_path.unlink()
except Exception:
pass
if temp_archive_path.exists():
temp_archive_path.unlink(missing_ok=True)
raise
else:
# Direct creation (non-atomic)
_create_archive_impl(archive_path, files, compression_level)
_create_archive_impl(archive_path_obj, files, compression_level)
def _create_archive_impl(
@@ -981,7 +969,7 @@ def test_archive(archive_path: str | Path, streaming: bool = False) -> bool:
if fileobj:
# Read the entire file to verify decompression
while True:
chunk = fileobj.read(8192)
chunk = fileobj.read(io.DEFAULT_BUFFER_SIZE)
if not chunk:
break
return True