main.py - master - Undeadguy/Tape_python

main.py
281 строка10 KB
Undeadguy
1.0
15 апр 2026, 22:40
15 апр 2026, 22:40556ac2a
import asyncio
import sys
import typing
from asyncio import Lock
from pathlib import Path
from typing import List, Optional

from rich import print
from typer import Argument, Option, Typer

from serializing_manager import SerManager

app = Typer()

stdout_aliases = ["STDOUT", "OUT", "1"]
stderr_aliases = ["STDERR", "ERR", "2"]

serialization_types = ["JSONL", "TOML"]

files_to_store: List[typing.Tuple[Path, int]] = []
files_lock: Lock = Lock()


@app.command()
def main(
    # Позиционные аргументы — файлы для объединения (как в cat)
    files: List[Path] = Argument(
        None,
        help="Файлы для объединения (если не указаны, используется директория)",
    ),
    # Опция для явного указания директории
    directory: Optional[Path] = Option(
        None,
        "--directory",
        "-d",
        help="Директория для рекурсивного сбора файлов (если не указаны конкретные файлы)",
    ),
    ignoring_files_str: List[str] = Option(
        [],
        "--ignoring-files",
        "-if",
        help="Игнорировать файлы с указанными именами (можно указать несколько раз)",
    ),
    ignoring_dirs_str: List[str] = Option(
        [],
        "--ignoring-dirs",
        "-id",
        help="Игнорировать директории с указанными именами (можно указать несколько раз)",
    ),
    ignoring_files_by_path: List[str] = Option(
        [],
        "-ifp",
        "--ignoring-files-path",
        help="Игнорировать файлы по полному пути (можно указать несколько раз)",
    ),
    ignoring_dirs_by_path: List[str] = Option(
        [],
        "-idp",
        "--ignoring-dirs-path",
        help="Игнорировать директории по полному пути (можно указать несколько раз)",
    ),
    store_path: typing.Optional[str] = Option(
        None,
        "--store-path",
        "-s",
        help="Куда сохранить результат: STDOUT/OUT/1 (stdout), STDERR/ERR/2 (stderr), или путь к файлу",
    ),
    serialize: str = Option(
        "TOML", "--serialize", help="Формат сериализации: TOML, JSONL (по умолчанию TOML)"
    ),
    create_file: bool = Option(
        False,
        "-c",
        "--create-file",
        help="Создавать файл назначения, если он не существует (только для store_path-файла)",
    ),
    ram_size_mb: Optional[int] = Option(None, "--max-ram-mb", "-ram"),
    size_limit_mb: Optional[int] = Option(None, "--size-limit"),
) -> None:

    # === ОПРЕДЕЛЯЕМ ИСТОЧНИК ФАЙЛОВ ===
    source_files = []

    if files:
        # Явно указаны файлы — используем их
        source_files = files
    #   print(f"[dim]Режим: объединение {len(files)} указанных файлов[/dim]")
    else:
        # Файлы не указаны — используем директорию
        start_dir = directory if directory else Path(".")
        start_dir = start_dir.resolve()

        if not start_dir.exists():
            print(
                f"[bold red] Ошибка! Директория {start_dir} не существует![/bold red]",
                file=sys.stderr,
            )
            sys.exit(1)
        elif not start_dir.is_dir():
            print(f"[bold red] Ошибка! {start_dir} — не директория![/bold red]", file=sys.stderr)
            sys.exit(2)

        # print(f"[dim]Режим: сбор всех файлов из директории {start_dir}[/dim]")

        # Валидация игнорируемых путей (как было)
        for file_path in ignoring_files_by_path:
            path = Path(file_path).resolve()
            if not path.exists():
                print(
                    f"[bold red] Ошибка! Путь {file_path} не существует![/bold red]",
                    file=sys.stderr,
                )
                sys.exit(1)
            elif path.is_dir():
                print(
                    f"[bold red] Ошибка! Путь {file_path} — директория, а ожидался файл![/bold red]",
                    file=sys.stderr,
                )
                sys.exit(2)

        for dir_path in ignoring_dirs_by_path:
            path = Path(dir_path).resolve()
            if not path.exists():
                print(
                    f"[bold red] Ошибка! Путь {dir_path} не существует![/bold red]", file=sys.stderr
                )
                sys.exit(1)
            elif path.is_file():
                print(
                    f"[bold red] Ошибка! Путь {dir_path} — файл, а ожидалась директория![/bold red]",
                    file=sys.stderr,
                )
                sys.exit(2)

        # Запускаем асинхронный сбор файлов из директории
        asyncio.run(
            async_collect_from_directory(
                start_dir,
                ignoring_files_str,
                ignoring_dirs_str,
                ignoring_files_by_path,
                ignoring_dirs_by_path,
                size_limit_mb,
            )
        )

    # Если есть явные файлы — обрабатываем их
    if source_files:
        # Валидация и сбор явных файлов
        for file_path in source_files:
            path = file_path.resolve()
            if not path.exists():
                print(
                    f"[bold red] Ошибка! Файл {file_path} не существует![/bold red]",
                    file=sys.stderr,
                )
                sys.exit(1)
            elif path.is_dir():
                print(
                    f"[bold red] Ошибка! {file_path} — директория, а ожидался файл![/bold red]",
                    file=sys.stderr,
                )
                sys.exit(2)

            size = path.stat().st_size
            files_to_store.append((path, size))

    # Если нет ни файлов, ни собранных из директории — ошибка
    if not files_to_store:
        print("[bold red] Ошибка! Нет файлов для обработки![/bold red]", file=sys.stderr)
        sys.exit(1)

    # === ДАЛЬШЕ СЕРИАЛИЗАЦИЯ ===
    if serialize.upper() not in serialization_types:
        print(
            f"[bold red] Ошибка! Сериализация {serialize} не поддерживается [/bold red]",
            file=sys.stderr,
        )
        sys.exit(1)

    # Определяем output
    output: typing.IO = sys.stdout
    if store_path:
        if store_path.upper() in stdout_aliases:
            output = sys.stdout
        elif store_path.upper() in stderr_aliases:
            output = sys.stderr
        else:
            path = Path(store_path).resolve()
            if not path.exists():
                if not create_file:
                    print("[bold red] Ошибка! Путь не существует [/bold red]", file=sys.stderr)
                    sys.exit(1)
                path.touch()
                output = path.open("w", encoding="utf-8")
            elif path.is_dir():
                print(
                    "[bold red] Ошибка! Вывод обязан быть либо в поток либо в файл! [/bold red]",
                    file=sys.stderr,
                )
                sys.exit(2)
            else:
                output = path.open("w", encoding="utf-8")

    # Запускаем сериализацию
    manager = SerManager(output, files_to_store, ram_size_mb, serialize)
    try:
        asyncio.run(manager.serialize_all())
    except Exception as e:
        print(f"[bold red] Ошибка на этапе сериализации! {e} [/bold red]", file=sys.stderr)
        sys.exit(1)

    print("[green bold]Команда выполнена успешно! [/green bold]")


async def async_collect_from_directory(
    start_directory: Path,
    ignoring_files_str: List[str],
    ignoring_dirs_str: List[str],
    ignoring_files_by_path: List[str],
    ignoring_dirs_by_path: List[str],
    size_limit_mb: Optional[int],
):
    """Асинхронный сбор файлов из директории"""
    await fs_walk(
        current_dir=start_directory,
        ignoring_dirs_by_path=ignoring_dirs_by_path,
        ignoring_dirs=ignoring_dirs_str,
        ignoring_files=ignoring_files_str,
        ignoring_files_by_path=ignoring_files_by_path,
        size_limit_mb=size_limit_mb,
    )


async def fs_walk(
    current_dir: Path,
    ignoring_files: List[str],
    ignoring_dirs: List[str],
    ignoring_files_by_path: List[str],
    ignoring_dirs_by_path: List[str],
    size_limit_mb: Optional[int],
):
    global files_to_store
    global files_lock

    selected_files: List[typing.Tuple[Path, int]] = []
    selected_dirs: List[Path] = []

    for entry in current_dir.iterdir():
        if entry.is_file():
            if len(ignoring_files) == 0 or entry.name not in ignoring_files:
                if len(ignoring_files_by_path) == 0 or str(entry) not in ignoring_files_by_path:
                    size = entry.stat().st_size
                    if (size_limit_mb is None) or (size / 1024**2 <= size_limit_mb):
                        selected_files.append((entry, size))
        elif entry.is_dir():
            if len(ignoring_dirs) == 0 or entry.name not in ignoring_dirs:
                if len(ignoring_dirs_by_path) == 0 or str(entry) not in ignoring_dirs_by_path:
                    selected_dirs.append(entry)

    async with files_lock:
        files_to_store.extend(selected_files)

    tasks = [
        fs_walk(
            subdir,
            ignoring_files,
            ignoring_dirs,
            ignoring_files_by_path,
            ignoring_dirs_by_path,
            size_limit_mb,
        )
        for subdir in selected_dirs
    ]

    if tasks:
        await asyncio.gather(*tasks)


if __name__ == "__main__":
    app()