file-splitter.py

#!/usr/bin/env python3
"""
File Splitter — Split large files into smaller chunks.

Usage:
    python file-splitter.py <file> --lines 1000
    python file-splitter.py <file> --size 10MB
    python file-splitter.py <file> --lines 500 --output chunks/

Options:
    --lines N        Split every N lines (text mode)
    --size SIZE      Split into chunks of max SIZE (e.g. 10MB, 500KB, 2GB)
    --output DIR     Output directory (default: <file>_chunks)
    --prefix PREFIX  Output file prefix (default: chunk)
    --suffix EXT     Output file extension (default: same as input)
    --no-header      Don't preserve first line as header in each chunk (for CSVs)
    --help           Show this help message and exit

Notes:
    - Use --lines for text files (preserves line integrity)
    - Use --size for binary files or byte-level splitting
"""

import os
import sys
import math
import argparse
from pathlib import Path


def parse_size(size_str: str) -> int:
    """Convert a human-readable size (e.g. '10MB', '500KB', '2GB') to bytes."""
    size_str = size_str.strip().upper()
    units = {"B": 1, "KB": 1024, "MB": 1024 ** 2, "GB": 1024 ** 3, "TB": 1024 ** 4}
    for unit, multiplier in units.items():
        if size_str.endswith(unit):
            try:
                num = float(size_str[: -len(unit)])
                return int(num * multiplier)
            except ValueError:
                break
    try:
        return int(size_str)
    except ValueError:
        print(f"Error: Invalid size format: '{size_str}'. Use e.g. 10MB, 500KB, 2GB.")
        sys.exit(1)


def split_by_lines(filepath: Path, lines_per_chunk: int,
                   out_dir: Path, prefix: str, suffix: str, keep_header: bool) -> int:
    """Split a text file by line count. Returns number of chunks."""
    chunk_num = 0
    header = None

    try:
        with open(filepath, "r", encoding="utf-8", errors="replace") as f:
            if keep_header:
                header = f.readline().rstrip("\n\r")

            while True:
                chunk_lines = []
                for _ in range(lines_per_chunk):
                    line = f.readline()
                    if not line:
                        break
                    chunk_lines.append(line.rstrip("\n\r"))

                if not chunk_lines:
                    break

                chunk_num += 1
                chunk_file = out_dir / f"{prefix}_{chunk_num:04d}{suffix}"
                with open(chunk_file, "w", encoding="utf-8") as cf:
                    if header:
                        cf.write(header + "\n")
                    for cl in chunk_lines:
                        cf.write(cl + "\n")
                print(f"  Created {chunk_file.name} ({len(chunk_lines)} lines)")

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

    return chunk_num


def split_by_size(filepath: Path, chunk_size: int,
                  out_dir: Path, prefix: str, suffix: str) -> int:
    """Split a file into fixed-size byte chunks. Returns number of chunks."""
    chunk_num = 0
    total_size = filepath.stat().st_size

    try:
        with open(filepath, "rb") as f:
            while True:
                data = f.read(chunk_size)
                if not data:
                    break
                chunk_num += 1
                chunk_file = out_dir / f"{prefix}_{chunk_num:04d}{suffix}"
                with open(chunk_file, "wb") as cf:
                    cf.write(data)
                pct = (chunk_num * chunk_size) / total_size * 100
                print(f"  Created {chunk_file.name} ({len(data)} bytes, {pct:.1f}%)")
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

    return chunk_num


def main():
    parser = argparse.ArgumentParser(
        description="Split large files into smaller chunks.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  python file-splitter.py huge-log.txt --lines 10000\n"
            "  python file-splitter.py video.mp4 --size 50MB\n"
            "  python file-splitter.py data.csv --lines 5000 --prefix data_chunk\n"
            "  python file-splitter.py data.csv --lines 5000 --no-header\n"
        ),
    )
    parser.add_argument("file", help="File to split")
    parser.add_argument("--lines", type=int, help="Lines per chunk (text mode)")
    parser.add_argument("--size", type=str, help="Chunk size (e.g. 10MB, 500KB)")
    parser.add_argument("--output", help="Output directory (default: <file>_chunks)")
    parser.add_argument("--prefix", default="chunk", help="Output file prefix (default: chunk)")
    parser.add_argument("--suffix", help="Output file extension (default: same as input)")
    parser.add_argument("--no-header", action="store_true", help="Don't repeat CSV header in each chunk")

    args = parser.parse_args()

    filepath = Path(args.file)
    if not filepath.is_file():
        print(f"Error: '{args.file}' is not a valid file.")
        sys.exit(1)

    if args.lines and args.size:
        print("Error: Use either --lines OR --size, not both.")
        sys.exit(1)
    if not args.lines and not args.size:
        print("Error: Specify either --lines or --size.")
        sys.exit(1)

    out_dir = Path(args.output) if args.output else filepath.parent / f"{filepath.name}_chunks"
    out_dir.mkdir(parents=True, exist_ok=True)

    suffix = args.suffix if args.suffix is not None else filepath.suffix

    file_size = filepath.stat().st_size
    print(f"File: {filepath} ({file_size:,} bytes)")

    if args.lines:
        print(f"Splitting every {args.lines} lines...")
        count = split_by_lines(filepath, args.lines, out_dir,
                               args.prefix, suffix, not args.no_header)
    else:
        chunk_size = parse_size(args.size)
        print(f"Splitting into {args.size} chunks ({chunk_size:,} bytes each)...")
        count = split_by_size(filepath, chunk_size, out_dir, args.prefix, suffix)

    print(f"\nDone: {count} chunk(s) created in {out_dir}")


if __name__ == "__main__":
    main()