conjurer/conjurer_librarian/cut_pieces.py

#!/usr/bin/env python3

from pathlib import Path
import argparse
import logging
import sys


# =========================
# Configuration constants
# =========================

MAX_CHUNK_SIZE_BYTES = 50 * 1024 * 1024  # 100 MB
OUTPUT_SUFFIX = "_chunk.txt"


# =========================
# Logging
# =========================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)

logger = logging.getLogger(__name__)


def get_default_output_dir(input_file: Path) -> Path:
    """
    Return default output directory placed next to the input file.

    Example:
        /data/big.txt -> /data/big/
    """
    return input_file.with_suffix("")


def write_chunk(output_dir: Path, chunk_index: int, lines: list[bytes], chunk_size: int) -> Path:
    """
    Write one chunk file and return its path.
    """
    output_file = output_dir / f"{chunk_index}{OUTPUT_SUFFIX}"

    with output_file.open("wb") as file:
        file.writelines(lines)

    actual_size = output_file.stat().st_size

    if actual_size != chunk_size:
        logger.warning(
            "Chunk size mismatch for %s: expected %d bytes, got %d bytes",
            output_file,
            chunk_size,
            actual_size,
        )

    return output_file


def split_txt_file(input_file: Path, output_dir: Path | None = None) -> None:
    """
    Split a text-like file into smaller files limited by MAX_CHUNK_SIZE_BYTES.

    The file is processed in binary mode, so invalid UTF-8 or mixed encodings
    do not break the split. Lines are never split.
    """
    if not input_file.exists():
        raise FileNotFoundError(f"Input file does not exist: {input_file}")

    if not input_file.is_file():
        raise ValueError(f"Input path is not a file: {input_file}")

    if MAX_CHUNK_SIZE_BYTES <= 0:
        raise ValueError("MAX_CHUNK_SIZE_BYTES must be greater than zero")

    if output_dir is None:
        output_dir = get_default_output_dir(input_file)

    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info("Input file: %s", input_file)
    logger.info("Output directory: %s", output_dir)
    logger.info("Max chunk size: %d bytes", MAX_CHUNK_SIZE_BYTES)

    chunk_index = 0
    current_lines: list[bytes] = []
    current_size = 0
    total_lines = 0
    total_bytes = 0

    with input_file.open("rb") as input_handle:
        for line_number, line in enumerate(input_handle, start=1):
            line_size = len(line)
            total_lines = line_number
            total_bytes += line_size

            if line_size > MAX_CHUNK_SIZE_BYTES:
                logger.warning(
                    "Line %d is larger than max chunk size: %d bytes > %d bytes. "
                    "It will be written as a separate chunk.",
                    line_number,
                    line_size,
                    MAX_CHUNK_SIZE_BYTES,
                )

            if current_lines and current_size + line_size > MAX_CHUNK_SIZE_BYTES:
                output_file = write_chunk(
                    output_dir=output_dir,
                    chunk_index=chunk_index,
                    lines=current_lines,
                    chunk_size=current_size,
                )

                logger.info(
                    "Written chunk %d: %s (%d bytes)",
                    chunk_index,
                    output_file,
                    current_size,
                )

                chunk_index += 1
                current_lines = []
                current_size = 0

            current_lines.append(line)
            current_size += line_size

    if current_lines:
        output_file = write_chunk(
            output_dir=output_dir,
            chunk_index=chunk_index,
            lines=current_lines,
            chunk_size=current_size,
        )

        logger.info(
            "Written chunk %d: %s (%d bytes)",
            chunk_index,
            output_file,
            current_size,
        )

        chunk_count = chunk_index + 1
    else:
        logger.info("Input file is empty. No chunks were created.")
        chunk_count = 0

    logger.info("Total lines processed: %d", total_lines)
    logger.info("Total bytes processed: %d", total_bytes)
    logger.info("Chunks created: %d", chunk_count)
    logger.info("Done.")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Split a large TXT-like file into smaller line-safe chunks."
    )

    parser.add_argument(
        "input_file",
        type=Path,
        help="Path to input TXT file.",
    )

    parser.add_argument(
        "-o",
        "--output-dir",
        type=Path,
        default=None,
        help=(
            "Optional output directory. "
            "By default, a directory named after the input file is created next to it."
        ),
    )

    return parser.parse_args()


def main() -> int:
    args = parse_args()

    try:
        split_txt_file(args.input_file, args.output_dir)
        return 0
    except Exception:
        logger.exception("Failed to split file")
        return 1


if __name__ == "__main__":
    sys.exit(main())