bx_py_utils/test_utils/zip_file_utils.py

import dataclasses
import datetime
import io
import os
import zipfile
from collections.abc import Generator
from unittest import mock
from zipfile import compressor_names

from bx_py_utils.dict_utils import dict_list2markdown
from bx_py_utils.test_utils.context_managers import MassContextManager
from bx_py_utils.test_utils.datetime import parse_dt


class FreezeZipFileDatetime(MassContextManager):
    """
    Context manager / decorator to freezes the modification time of files written to a zip file.
    (e.g.: Helps to create reproducible openpyxl files)
    """

    def __init__(self, dt_str: str = '2024-12-24T00:00:00+00:00'):
        dt = parse_dt(dt_str)
        timestamp = dt.timestamp()
        real_write = zipfile.ZipFile.write

        def _mock_write(zf_self, filename, *args, **kwargs):
            os.utime(filename, (timestamp, timestamp))
            return real_write(zf_self, filename, *args, **kwargs)

        self.mocks = (mock.patch('zipfile.ZipFile.write', _mock_write),)


@dataclasses.dataclass
class ZipFileInfo:
    file_size: int
    compress_size: int
    compress_ratio: int
    compress_type: str
    modification_date: str  # Modification date in ISO format
    crc32: str  # CRC-32 of the uncompressed file
    file_name: str


def zip_info(file) -> Generator[ZipFileInfo, None, None]:
    """
    Generates similar information than `unzip -v`: Yields ZipFileInfo for each file in the zip file.
    It could be used to create a snapshot of zip file content, e.g.:

        assert_py_snapshot(got=list(zip_info(zip_file)))

    Hint: xlsx (e.g.: generated by openpyxl) are also normal zip files, if you need to test them ;)
    Hint2: Can be easy combined with bx_py_utils.dict_utils.dict_list2markdown() and our snapshot functions ;)
    """
    if isinstance(file, bytes):
        file = io.BytesIO(file)
    with zipfile.ZipFile(file, 'r') as zf:
        for info in zf.infolist():
            dt = datetime.datetime(*info.date_time)
            modification_date = dt.astimezone(datetime.timezone.utc)
            yield ZipFileInfo(
                file_size=info.file_size,
                compress_size=info.compress_size,
                compress_ratio=info.compress_size * 100 // info.file_size,
                compress_type=compressor_names.get(info.compress_type, info.compress_type),
                modification_date=modification_date.isoformat(timespec='seconds'),
                crc32=f'{info.CRC:08X}',
                file_name=info.filename,
            )


def zip_info_markdown(file):
    """
    Generates a markdown representation of the zip file content. Similar to `unzip -v` output.
    """
    return '\n'.join(dict_list2markdown(dataclasses.asdict(file_info) for file_info in zip_info(file)))