""" Tools for creating ``.conda``-format archives. Uses ``tempfile.SpooledTemporaryFile`` to buffer ``pkg-*.tar`` and ``info-*.tar``, then compress directly into an open `ZipFile` at the end. `SpooledTemporaryFile` buffers the first 10MB of the package and its metadata in memory, but writes out to disk for larger packages. Uses more disk space than ``conda-package-handling`` (temporary uncompressed tarballs of the package contents) but accepts streams instead of just files-on-the-filesystem. """ from __future__ import annotations import json import shutil import tarfile import tempfile import zipfile from collections.abc import Iterator from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING try: import compression.zstd as zstd except ImportError: import backports.zstd as zstd if TYPE_CHECKING: from collections.abc import Callable from typing import BinaryIO, Protocol class LegacyCompressor(Protocol): def stream_writer( self, writer: BinaryIO, *, size: int, closefd: bool, ) -> BinaryIO: ... LegacyCompressorOrFactory = LegacyCompressor | Callable[[], LegacyCompressor] # increase to reduce speed and increase compression (levels above 19 use much # more memory) ZSTD_COMPRESS_LEVEL = 19 # increase for greater speed ZSTD_COMPRESS_THREADS = 1 CONDA_PACKAGE_FORMAT_VERSION = 2 # Account for growth from "2 GB of /dev/urandom" to not exceed ZIP64_LIMIT after # compression CONDA_ZIP64_LIMIT = zipfile.ZIP64_LIMIT - (1 << 18) - 1 def anonymize(tarinfo: tarfile.TarInfo): """ Pass to ``tarfile.add(..., filter=anonymize)`` to anonymize uid/gid. Does not anonymize mtime or any other field. """ tarinfo.uid = tarinfo.gid = 0 tarinfo.uname = tarinfo.gname = "" return tarinfo class CondaTarFile(tarfile.TarFile): """ Subclass of :external+python:py:class:`tarfile.TarFile` that adds members to a second ``info`` tar if they match ``is_info(name)``. Create this with ``conda_builder(...)`` which sets up the component archives, then wraps them into a ``.conda`` on exit. Only useful for creating, not extracting ``.conda``. """ info_tar: tarfile.TarFile is_info: Callable def __init__( self, *args, info_tar: tarfile.TarFile, is_info=lambda name: name.startswith("info/"), **kwargs, ): super().__init__(*args, **kwargs) self.info_tar = info_tar self.is_info = is_info def addfile(self, tarinfo, fileobj=None): """ Add the TarInfo object ``tarinfo`` to the archive. If ``fileobj`` is given, it should be a binary file, and tarinfo.size bytes are read from it and added to the archive. You can create TarInfo objects directly, or by using ``gettarinfo()``. If ``self.is_info(tarinfo.name)`` returns ``True``, add ``tarinfo`` to ``self.info_tar`` instead. """ if self.is_info(tarinfo.name): return self.info_tar.addfile(tarinfo, fileobj=fileobj) else: return super().addfile(tarinfo, fileobj) class _ZstdFile(zstd.ZstdFile): """ ZstdFile adding pledged_input_size. For the older zstandard binding, a pledged size was important for decompression memory usage. Check if that is still true in compression.zstd. """ def __init__(self, *args, pledged_input_size=0, **kwargs): super().__init__(*args, **kwargs) self._compressor.set_pledged_input_size(pledged_input_size) def _open_writer( output, *, pledged_input_size: int, compressor: LegacyCompressorOrFactory | None, compression_level: int | None = None, compression_threads: int | None = None, ): if compressor is not None: if compression_level is not None or compression_threads is not None: msg = "`compressor` overrides `compression_level` and `compression_threads`" raise ValueError(msg) legacy_compressor = compressor() if callable(compressor) else compressor if not hasattr(legacy_compressor, "stream_writer"): msg = "compressor must provide stream_writer(...)" raise TypeError(msg) return legacy_compressor.stream_writer( output, size=pledged_input_size, closefd=False, ) if compression_level is None: compression_level = ZSTD_COMPRESS_LEVEL if compression_threads is None: compression_threads = ZSTD_COMPRESS_THREADS return _ZstdFile( output, mode="w", pledged_input_size=pledged_input_size, options={ zstd.CompressionParameter.compression_level: compression_level, zstd.CompressionParameter.nb_workers: compression_threads, }, ) @contextmanager def conda_builder( stem, path, *, compressor: LegacyCompressorOrFactory | None = None, compression_level: int | None = None, compression_threads: int | None = None, is_info: Callable[[str], bool] = lambda filename: filename.startswith("info/"), encoding="utf-8", ) -> Iterator[CondaTarFile]: """ Produce a ``TarFile`` subclass used to build a ``.conda`` package. The subclass delegates ``addfile()`` to the ``info-`` component when ``is_info`` returns True. When the context manager exits, ``{path}/{stem}.conda`` is written with the component tar archives. Args: stem: output filename without extension path: destination path for transmuted .conda package. compressor: Legacy ``zstandard`` compressor object (or factory returning one) with ``stream_writer(...)``. Mutually exclusive with ``compression_level`` and ``compression_threads``. compression_level: zstd compression level for ``compression.zstd`` or ``backports.zstd`` code path. Defaults to ``ZSTD_COMPRESS_LEVEL`` if not specified and ``compressor`` is None. compression_threads: Number of zstd worker threads for ``compression.zstd`` or ``backports.zstd`` code path. Defaults to ``ZSTD_COMPRESS_THREADS`` if not specified and ``compressor`` is None. encoding: passed to TarFile constructor. Keep default "utf-8" for valid .conda. Yields: ``CondaTarFile`` """ output_path = Path(path, f"{stem}.conda") with ( tempfile.SpooledTemporaryFile() as info_file, tempfile.SpooledTemporaryFile() as pkg_file, ): with ( tarfile.TarFile(fileobj=info_file, mode="w", encoding=encoding) as info_tar, CondaTarFile( fileobj=pkg_file, mode="w", info_tar=info_tar, is_info=is_info, encoding=encoding, ) as pkg_tar, ): # If we wanted to compress these at a low setting to save temporary # space, we could insert a file object that counts bytes written in # front of a zstd (level between 1..3) compressor. yield pkg_tar info_tar.close() pkg_tar.close() info_size = info_file.tell() pkg_size = pkg_file.tell() info_file.seek(0) pkg_file.seek(0) with zipfile.ZipFile( output_path, "x", # x to not append to existing compresslevel=zipfile.ZIP_STORED, ) as conda_file: pkg_metadata = {"conda_pkg_format_version": CONDA_PACKAGE_FORMAT_VERSION} conda_file.writestr("metadata.json", json.dumps(pkg_metadata)) with ( conda_file.open( f"pkg-{stem}.tar.zst", "w", force_zip64=(pkg_size > CONDA_ZIP64_LIMIT), ) as pkg_file_zip, _open_writer( pkg_file_zip, pledged_input_size=pkg_size, compressor=compressor, compression_level=compression_level, compression_threads=compression_threads, ) as pkg_stream, ): shutil.copyfileobj(pkg_file._file, pkg_stream) with ( conda_file.open( f"info-{stem}.tar.zst", "w", force_zip64=(info_size > CONDA_ZIP64_LIMIT), ) as info_file_zip, _open_writer( info_file_zip, pledged_input_size=info_size, compressor=compressor, compression_level=compression_level, compression_threads=compression_threads, ) as info_stream, ): shutil.copyfileobj(info_file._file, info_stream)