archiver.py - main - 2FIVE/Laba1_python_arhivator

archiver.py
652 строки26 KB
2FIVE
update archiver.py
03 янв 2026, 12:46
03 янв 2026, 12:468dc94d7
"""
Консольный архиватор/распаковщик с поддержкой zstd и bz2
Использует только стандартные библиотеки Python 3.14
"""

import argparse
import os
import sys
import time
import tarfile
import bz2
import io
import itertools
import math
from pathlib import Path
from typing import Optional, Tuple, Callable, Any
from dataclasses import dataclass
from enum import Enum

# Проверка версии Python
import sys
print(f"Python версия: {sys.version}")
print(f"Python путь: {sys.executable}")

# Проверяем, что это Python 3.14+
if sys.version_info < (3, 14):
    print("ОШИБКА: Требуется Python 3.14 или новее")
    print(f"   Текущая версия: {sys.version_info.major}.{sys.version_info.minor}")
    print("   Установите Python 3.14 или используйте явный путь:")
    print(r"   C:\Users\Илья\AppData\Local\Programs\Python\Python314\python.exe") 
    sys.exit(1) # укажите свой путь если потребуется, либо удалите или закомментируйте данный участок кода
    
# Настройка кодировки для корректного отображения русского текста
if sys.platform == "win32":
    # Для Windows устанавливаем кодировку UTF-8
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
    os.environ['PYTHONIOENCODING'] = 'utf-8'

# Прямой импорт zstd согласно PEP 784
try:
    from compression import zstd
    print("[INFO] Используется compression.zstd (PEP 784)")
    HAS_ZSTD = True
except ImportError as e:
    print(f"[ОШИБКА] Не удалось импортировать compression.zstd: {e}")
    print("[INFO] Убедитесь, что используете Python 3.14+")
    HAS_ZSTD = False
    sys.exit(1)


class CompressionType(Enum):
    """Типы поддерживаемого сжатия"""
    ZSTD = 'zstd'
    BZ2 = 'bz2'
    TAR = 'tar'
    TAR_ZSTD = 'tar.zst'
    TAR_BZ2 = 'tar.bz2'


class ProgressIndicator:
    """Индикатор прогресса для консоли без многопоточности"""
    
    def __init__(self, total: int, enabled: bool = True):
        self.total = total
        self.current = 0
        self.enabled = enabled
        self.start_time = time.time()
        # Анимация прогресс-бара
        self.animation_cycle = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'])
        self.last_update = 0
        self.width = 50
        
    def update(self, value: int) -> None:
        """Обновить значение прогресса"""
        self.current = min(value, self.total)
        
        if not self.enabled:
            return
            
        current_time = time.time()
        if current_time - self.last_update < 0.1:
            return
            
        self.last_update = current_time
        self._display()
        
    def _display(self) -> None:
        """Отобразить текущий прогресс"""
        if self.total <= 0:
            return
            
        progress = min(self.current / self.total, 1.0)
        filled_width = int(self.width * progress)
        bar = '█' * filled_width + '░' * (self.width - filled_width)
        percent = progress * 100
        
        elapsed = time.time() - self.start_time
        if progress > 0:
            estimated_total = elapsed / progress
            remaining = estimated_total - elapsed
            time_str = f" {remaining:.1f}с осталось"
        else:
            time_str = ""
            
        anim_char = next(self.animation_cycle)
        sys.stdout.write(f'\r{anim_char} [{bar}] {percent:5.1f}%{time_str}')
        sys.stdout.flush()
        
    def finish(self) -> None:
        """Завершить отображение прогресса"""
        if not self.enabled:
            return
            
        bar = '█' * self.width
        elapsed = time.time() - self.start_time
        
        sys.stdout.write(f'\r✓ [{bar}] 100.0% Завершено за {elapsed:.1f}с\n')
        sys.stdout.flush()


@dataclass
class ArchiveInfo:
    """Информация об архиве"""
    compression_type: CompressionType
    is_tar_archive: bool
    base_name: str


class Archiver:
    """Класс для архивации и распаковки с поддержкой zstd и bz2"""
    
    def __init__(self, benchmark: bool = False, progress: bool = False):
        self.benchmark = benchmark
        self.progress = progress
        
    def _detect_archive_type(self, filename: Path) -> ArchiveInfo:
        """Определить тип архива по расширению файла"""
        name = filename.name.lower()
        
        if name.endswith('.tar.zst'):
            return ArchiveInfo(
                compression_type=CompressionType.TAR_ZSTD,
                is_tar_archive=True,
                base_name=name[:-8]
            )
        elif name.endswith('.tar.bz2'):
            return ArchiveInfo(
                compression_type=CompressionType.TAR_BZ2,
                is_tar_archive=True,
                base_name=name[:-8]
            )
        elif name.endswith('.zst'):
            return ArchiveInfo(
                compression_type=CompressionType.ZSTD,
                is_tar_archive=False,
                base_name=name[:-4]
            )
        elif name.endswith('.bz2'):
            return ArchiveInfo(
                compression_type=CompressionType.BZ2,
                is_tar_archive=False,
                base_name=name[:-4]
            )
        elif name.endswith('.tar'):
            return ArchiveInfo(
                compression_type=CompressionType.TAR,
                is_tar_archive=True,
                base_name=name[:-4]
            )
        else:
            raise ValueError(f"[ОШИБКА] Неподдерживаемый формат архива: {filename}")
    
    def _compress_zstd(self, source: Path, target: Path, 
                      progress: Optional[ProgressIndicator] = None) -> bool:
        """Сжать файл с использованием zstd (как в PEP 784)"""
        try:
            chunk_size = 64 * 1024  # 64KB
            total_size = source.stat().st_size
            processed = 0
            
            with open(source, 'rb') as f_in:
                # Собираем данные с отслеживанием прогресса
                chunks = []
                while True:
                    chunk = f_in.read(chunk_size)
                    if not chunk:
                        break
                    chunks.append(chunk)
                    processed += len(chunk)
                    
                    if progress and total_size > 0:
                        progress.update(int((processed / total_size) * 100))
                
                # Объединяем данные
                data = b''.join(chunks)
                
                # Сжимаем данные как в примере PEP 784
                compressed = zstd.compress(data, level=3)
                
                # Записываем результат
                with open(target, 'wb') as f_out:
                    f_out.write(compressed)
            
            return True
            
        except Exception as e:
            print(f"[ОШИБКА] При сжатии zstd: {e}")
            return False
    
    def _decompress_zstd(self, source: Path, target: Path,
                        progress: Optional[ProgressIndicator] = None) -> bool:
        """Распаковать ZSTD файл (как в PEP 784)"""
        try:
            chunk_size = 64 * 1024
            total_size = source.stat().st_size
            processed = 0
            
            with open(source, 'rb') as f_in:
                # Читаем сжатые данные
                compressed_data = f_in.read()
                processed = total_size
                if progress and total_size > 0:
                    progress.update(100)
                
                # Распаковываем
                decompressed = zstd.decompress(compressed_data)
                
                # Записываем результат
                with open(target, 'wb') as f_out:
                    f_out.write(decompressed)
            
            return True
            
        except Exception as e:
            print(f"[ОШИБКА] При распаковке zstd: {e}")
            return False
    
    def _compress_bz2(self, source: Path, target: Path,
                     progress: Optional[ProgressIndicator] = None) -> bool:
        """Сжать файл с использованием bz2"""
        try:
            chunk_size = 64 * 1024
            total_size = source.stat().st_size
            processed = 0
            
            with open(source, 'rb') as f_in:
                with bz2.open(target, 'wb', compresslevel=9) as f_out:
                    while True:
                        chunk = f_in.read(chunk_size)
                        if not chunk:
                            break
                        f_out.write(chunk)
                        processed += len(chunk)
                        
                        if progress and total_size > 0:
                            progress.update(int((processed / total_size) * 100))
            
            return True
            
        except Exception as e:
            print(f"[ОШИБКА] При сжатии bz2: {e}")
            return False
    
    def _decompress_bz2(self, source: Path, target: Path,
                       progress: Optional[ProgressIndicator] = None) -> bool:
        """Распаковать bz2 файл"""
        try:
            chunk_size = 64 * 1024
            total_size = source.stat().st_size
            processed = 0
            
            with bz2.open(source, 'rb') as f_in:
                with open(target, 'wb') as f_out:
                    while True:
                        chunk = f_in.read(chunk_size)
                        if not chunk:
                            break
                        f_out.write(chunk)
                        processed += len(chunk)
                        
                        if progress and total_size > 0:
                            progress.update(int((processed / total_size) * 100))
            
            return True
            
        except Exception as e:
            print(f"[ОШИБКА] При распаковке bz2: {e}")
            return False
    
    def _create_tar_with_progress(self, source: Path, target: Path,
                                progress: Optional[ProgressIndicator] = None) -> bool:
        """Создать tar архив с отслеживанием прогресса"""
        try:
            file_count = 0
            if source.is_dir():
                file_count = sum(1 for _ in source.rglob('*') if _.is_file())
            else:
                file_count = 1
            
            current_file = 0
            
            with tarfile.open(target, 'w') as tar:
                if source.is_dir():
                    for root, dirs, files in os.walk(source):
                        for file in files:
                            filepath = os.path.join(root, file)
                            arcname = os.path.relpath(filepath, source.parent)
                            tar.add(filepath, arcname=arcname)
                            current_file += 1
                            if progress and file_count > 0:
                                progress.update(int((current_file / file_count) * 100))
                else:
                    tar.add(source, arcname=source.name)
                    current_file = 1
                    if progress:
                        progress.update(100)
            
            return True
            
        except Exception as e:
            print(f"[ОШИБКА] При создании tar архива: {e}")
            return False
    
    def _extract_tar_with_progress(self, source: Path, target_dir: Path,
                                 progress: Optional[ProgressIndicator] = None) -> bool:
        """Распаковать tar архив с отслеживанием прогресса"""
        try:
            with tarfile.open(source, 'r') as tar:
                members = tar.getmembers()
                if progress:
                    progress.total = len(members)
                
                for i, member in enumerate(members):
                    tar.extract(member, target_dir)
                    if progress:
                        progress.update(i + 1)
            
            return True
            
        except Exception as e:
            print(f"[ОШИБКА] При распаковке tar архива: {e}")
            return False
    
    def _calculate_directory_size(self, path: Path) -> int:
        """Вычислить общий размер директории"""
        if path.is_file():
            return path.stat().st_size
        
        total = 0
        for item in path.rglob('*'):
            if item.is_file():
                total += item.stat().st_size
        return total
    
    def _format_size(self, size_bytes: int) -> str:
        """Форматировать размер в читаемый вид"""
        if size_bytes == 0:
            return "0 Б"
        
        size_names = ("Б", "КБ", "МБ", "ГБ", "ТБ")
        i = int(math.floor(math.log(size_bytes, 1024)))
        p = math.pow(1024, i)
        s = round(size_bytes / p, 2)
        return f"{s} {size_names[i]}"
    
    def compress(self, input_path: str, output_file: str) -> bool:
        """Архивировать файл или директорию"""
        start_time = time.time() if self.benchmark else None
        
        try:
            source = Path(input_path)
            target = Path(output_file)
            
            if not source.exists():
                print(f"[ОШИБКА] Путь '{source}' не существует")
                return False
            
            archive_info = self._detect_archive_type(target)
            
            print(f" Архивация: {source.name} -> {target.name} ({archive_info.compression_type.value})")
            
            progress = None
            if self.progress:
                progress = ProgressIndicator(100, enabled=True)
            
            # Решаем, нужно ли создавать tar
            needs_tar = source.is_dir() or archive_info.is_tar_archive
            
            success = False
            
            if needs_tar:
                # Для директорий всегда создаем tar
                temp_tar = target.with_suffix('.tar')
                
                print(f" Создание tar архива...")
                if progress:
                    progress.update(0)
                
                if not self._create_tar_with_progress(source, temp_tar, progress):
                    return False
                
                # Теперь сжимаем tar файл
                if archive_info.compression_type in [CompressionType.TAR_ZSTD, CompressionType.ZSTD]:
                    print(f" Сжатие zstd...")
                    if progress:
                        progress = ProgressIndicator(100, enabled=self.progress)
                    success = self._compress_zstd(temp_tar, target, progress)
                elif archive_info.compression_type in [CompressionType.TAR_BZ2, CompressionType.BZ2]:
                    print(f" Сжатие bz2...")
                    if progress:
                        progress = ProgressIndicator(100, enabled=self.progress)
                    success = self._compress_bz2(temp_tar, target, progress)
                elif archive_info.compression_type == CompressionType.TAR:
                    temp_tar.rename(target)
                    success = True
                
                # Удаляем временный tar файл
                if temp_tar.exists() and temp_tar != target:
                    temp_tar.unlink()
                    
            else:
                # Одиночный файл, сжимаем напрямую
                print(f" Сжатие {archive_info.compression_type.value}...")
                if progress:
                    progress.update(0)
                
                if archive_info.compression_type == CompressionType.ZSTD:
                    success = self._compress_zstd(source, target, progress)
                else:  # BZ2
                    success = self._compress_bz2(source, target, progress)
            
            if progress:
                progress.finish()
            
            if success:
                input_size = self._calculate_directory_size(source)
                output_size = target.stat().st_size
                compression_ratio = (1 - output_size / input_size) * 100 if input_size > 0 else 0
                
                print(f" Статистика:")
                print(f"   Исходный размер: {self._format_size(input_size)}")
                print(f"   Размер архива:   {self._format_size(output_size)}")
                print(f"   Коэффициент сжатия: {compression_ratio:.1f}%")
                
                if self.benchmark and start_time:
                    elapsed = time.time() - start_time
                    print(f"⏱  Время выполнения: {elapsed:.2f} секунд")
                
                print(f" Архивация успешно завершена!")
            
            return success
            
        except Exception as e:
            print(f"[ОШИБКА] При архивации: {e}")
            return False
    
    def decompress(self, input_file: str, output_dir: Optional[str] = None) -> bool:
        """Распаковать архив"""
        start_time = time.time() if self.benchmark else None
        
        try:
            source = Path(input_file)
            
            if not source.exists():
                print(f"[ОШИБКА] Файл '{source}' не существует")
                return False
            
            archive_info = self._detect_archive_type(source)
            
            if output_dir:
                target_dir = Path(output_dir)
            else:
                target_dir = Path.cwd() / archive_info.base_name
            
            print(f" Распаковка: {source.name} -> {target_dir}/ ({archive_info.compression_type.value})")
            
            progress = None
            if self.progress:
                progress = ProgressIndicator(100, enabled=True)
            
            success = False
            
            if archive_info.compression_type in [CompressionType.TAR_ZSTD, CompressionType.TAR_BZ2]:
                # Комбинированный архив
                print(f" Распаковка сжатия...")
                
                temp_tar = target_dir.with_suffix('.tar')
                
                if progress:
                    progress.update(0)
                
                if archive_info.compression_type == CompressionType.TAR_ZSTD:
                    decompress_ok = self._decompress_zstd(source, temp_tar, progress)
                else:  # TAR_BZ2
                    decompress_ok = self._decompress_bz2(source, temp_tar, progress)
                
                if decompress_ok and temp_tar.exists():
                    print(f" Распаковка tar архива...")
                    if progress:
                        progress = ProgressIndicator(100, enabled=self.progress)
                    target_dir.mkdir(parents=True, exist_ok=True)
                    success = self._extract_tar_with_progress(temp_tar, target_dir, progress)
                    temp_tar.unlink()
                else:
                    success = False
                    
            elif archive_info.compression_type == CompressionType.TAR:
                # Простой tar архив
                print(f" Распаковка tar архива...")
                if progress:
                    progress.update(0)
                target_dir.mkdir(parents=True, exist_ok=True)
                success = self._extract_tar_with_progress(source, target_dir, progress)
                
            elif archive_info.compression_type in [CompressionType.ZSTD, CompressionType.BZ2]:
                # Одиночный сжатый файл
                print(f" Распаковка файла...")
                if progress:
                    progress.update(0)
                
                target_dir.mkdir(parents=True, exist_ok=True)
                
                # Выходной файл
                output_file = target_dir / archive_info.base_name
                
                if archive_info.compression_type == CompressionType.ZSTD:
                    success = self._decompress_zstd(source, output_file, progress)
                else:  # BZ2
                    success = self._decompress_bz2(source, output_file, progress)
            else:
                print(f"[ОШИБКА] Неподдерживаемый тип архива: {archive_info.compression_type}")
                success = False
            
            if progress:
                progress.finish()
            
            if success:
                print(f" Распаковка успешно завершена в {target_dir}/")
                
                if self.benchmark and start_time:
                    elapsed = time.time() - start_time
                    print(f"⏱ Время выполнения: {elapsed:.2f} секунд")
            
            return success
            
        except Exception as e:
            print(f"[ОШИБКА] При распаковке: {e}")
            return False


def main():
    """Главная функция утилиты"""
    parser = argparse.ArgumentParser(
        description='Консольный архиватор/распаковщик с поддержкой zstd и bz2',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Примеры использования:
  Архивация файла:    python archiver.py compress file.txt archive.zst
  Архивация файла:    python archiver.py compress file.txt archive.bz2
  Архивация папки:    python archiver.py compress my_folder archive.tar.zst
  Архивация папки:    python archiver.py compress my_folder archive.tar.bz2
  Распаковка:         python archiver.py decompress archive.tar.zst
  Распаковка в папку: python archiver.py decompress archive.bz2 output_folder
  С benchmark:        python archiver.py compress file.txt archive.zst --benchmark
  С прогресс-баром:   python archiver.py compress file.txt archive.zst --progress

Поддерживаемые форматы:
  .zst       - zstd сжатый файл (реальный Zstandard алгоритм, PEP 784)
  .bz2       - bz2 сжатый файл
  .tar       - tar архив
  .tar.zst   - tar архив + zstd сжатие
  .tar.bz2   - tar архив + bz2 сжатие

Требования: Python 3.14 или новее
        '''
    )
    
    subparsers = parser.add_subparsers(
        dest='command',
        help='Команда',
        required=True
    )
    
    compress_parser = subparsers.add_parser(
        'compress',
        help='Архивация файла или директории'
    )
    compress_parser.add_argument(
        'input',
        help='Входной файл или директория'
    )
    compress_parser.add_argument(
        'output',
        help='Выходной архив (.zst, .bz2, .tar.zst, .tar.bz2, .tar)'
    )
    compress_parser.add_argument(
        '--benchmark',
        action='store_true',
        help='Показать время выполнения'
    )
    compress_parser.add_argument(
        '--progress',
        action='store_true',
        help='Показать прогресс-бар'
    )
    
    decompress_parser = subparsers.add_parser(
        'decompress',
        help='Распаковка архива'
    )
    decompress_parser.add_argument(
        'input',
        help='Входной архив'
    )
    decompress_parser.add_argument(
        'output',
        nargs='?',
        help='Выходная директория (опционально)'
    )
    decompress_parser.add_argument(
        '--benchmark',
        action='store_true',
        help='Показать время выполнения'
    )
    decompress_parser.add_argument(
        '--progress',
        action='store_true',
        help='Показать прогресс-бар'
    )
    
    args = parser.parse_args()
    
    archiver = Archiver(
        benchmark=args.benchmark,
        progress=args.progress
    )
    
    if args.command == 'compress':
        success = archiver.compress(args.input, args.output)
    elif args.command == 'decompress':
        success = archiver.decompress(args.input, args.output)
    else:
        print("[ОШИБКА] Неизвестная команда")
        success = False
    
    sys.exit(0 if success else 1)


if __name__ == '__main__':
    main()