main.py - master - nedeadinside/backup_Task1

main.py
188 строк6 KB
nedeadinside
Init backup commit
14 янв 2026, 12:04
14 янв 2026, 12:049eeacb8
import bz2
import tarfile
from pathlib import Path
from compression import zstd

import sys
import time
import argparse
from typing import Literal

_CHUNK_SIZE = 1024 * 1024


def compress_file(
    source_path: Path, target_path: Path, algorithm: Literal["zstd", "bz2"]
) -> None:
    match algorithm:
        case "zstd":
            czstd = zstd.ZstdCompressor()
            target_path.parent.mkdir(parents=True, exist_ok=True)
            with source_path.open("rb") as f_in, target_path.open("wb") as f_out:
                for chunk in iter(lambda: f_in.read(_CHUNK_SIZE), b""):
                    out = czstd.compress(chunk, mode=zstd.ZstdCompressor.CONTINUE)
                    if out:
                        f_out.write(out)
                tail = czstd.flush(mode=zstd.ZstdCompressor.FLUSH_FRAME)
                if tail:
                    f_out.write(tail)
        case "bz2":
            target_path.parent.mkdir(parents=True, exist_ok=True)
            with (
                source_path.open("rb") as f_in,
                bz2.open(target_path, "wb", compresslevel=9) as f_out,
            ):
                for chunk in iter(lambda: f_in.read(_CHUNK_SIZE), b""):
                    f_out.write(chunk)


def decompress_file(
    source_path: Path, target_path: Path, algorithm: Literal["zstd", "bz2"]
) -> None:
    match algorithm:
        case "zstd":
            dzstd = zstd.ZstdDecompressor()
            target_path.parent.mkdir(parents=True, exist_ok=True)
            with source_path.open("rb") as f_in, target_path.open("wb") as f_out:
                for chunk in iter(lambda: f_in.read(_CHUNK_SIZE), b""):
                    out = dzstd.decompress(chunk)
                    if out:
                        f_out.write(out)
                    if dzstd.eof:
                        break
        case "bz2":
            with bz2.open(source_path, "rb") as f_in, target_path.open("wb") as f_out:
                for chunk in iter(lambda: f_in.read(_CHUNK_SIZE), b""):
                    f_out.write(chunk)


def get_compression_algorithm(file_path: Path) -> Literal["zstd", "bz2"]:
    suffix = file_path.suffix.lower()
    match suffix:
        case ".zst":
            return "zstd"
        case ".bz2":
            return "bz2"
        case _:
            raise ValueError(f"Неподдерживаемое расширение: {suffix}")


def archive_path(
    source_path: Path, target_path: Path, algorithm: Literal["zstd", "bz2"]
) -> None:
    if source_path.is_dir():
        tar_path = target_path.with_suffix(".tar")
        with tarfile.open(tar_path, "w") as tar:
            tar.add(source_path, arcname=source_path.name)
        try:
            compress_file(tar_path, target_path, algorithm)
        finally:
            if tar_path.exists():
                tar_path.unlink()
    else:
        compress_file(source_path, target_path, algorithm)


def extract_path(
    source_path: Path, target_path: Path, algorithm: Literal["zstd", "bz2"]
) -> None:
    temp_tar = target_path.parent / f"{target_path.name}.tmp.tar"
    temp_tar.parent.mkdir(parents=True, exist_ok=True)

    try:
        decompress_file(source_path, temp_tar, algorithm)
        if tarfile.is_tarfile(temp_tar):
            target_path.mkdir(parents=True, exist_ok=True)
            with tarfile.open(temp_tar, "r:") as tar:
                tar.extractall(path=target_path, filter="data")
        else:
            if target_path.exists() and target_path.is_dir():
                dst = target_path / source_path.stem
            else:
                dst = target_path
            try:
                temp_tar.replace(dst)
            except Exception:
                with open(temp_tar, "rb") as src, open(dst, "wb") as dst_f:
                    dst_f.write(src.read())
    finally:
        if temp_tar.exists():
            try:
                temp_tar.unlink()
            except Exception:
                pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Утилита архивации и распаковки с алгоритмами zstd и bz2",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Примеры использования:\n"
            "\tСжать файл в .zst: \tpython main.py -c source.txt out.zst\n"
            "\tСжать папку в .bz2:\tpython main.py --compress mydir archive.tar.bz2\n"
            "\tРаспаковать .zst:  \tpython main.py -d archive.tar.zst extracted_dir\n"
            "\tПоказать время:    \tpython main.py -c in.txt out.zst --verbose\n"
        ),
    )

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "-c",
        "--compress",
        action="store_true",
        help="Флаг сжатия файла или папки",
    )
    group.add_argument(
        "-d",
        "--decompress",
        action="store_true",
        help="Флаг распаковки файла или папки",
    )

    parser.add_argument(
        "src_path",
        type=Path,
        help="Путь к источнику",
    )
    parser.add_argument(
        "dst_path",
        type=Path,
        help="Путь назначения",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Выводить время выполнения",
    )

    args = parser.parse_args()

    if not args.src_path.exists():
        parser.error(f"Источник не найден: {args.src_path}")

    start = time.perf_counter() if args.verbose else None

    try:
        if args.compress:
            try:
                algorithm = get_compression_algorithm(args.dst_path)
            except ValueError as e:
                parser.error(str(e))
            archive_path(args.src_path, args.dst_path, algorithm)
        elif args.decompress:
            try:
                algorithm = get_compression_algorithm(args.src_path)
            except ValueError as e:
                parser.error(str(e))
            extract_path(args.src_path, args.dst_path, algorithm)
        else:
            parser.print_help()
            sys.exit(2)
    except Exception as exc:
        print(f"Ошибка: {exc}")
        sys.exit(1)
    finally:
        if start is not None:
            duration = time.perf_counter() - start
            print(f"Время выполнения: {duration:.3f} сек.")