split_excel.py

"""
split_excel.py — Split one spreadsheet into many files.

Two modes:
  * --by COLUMN : write one file per unique value in that column
                  (e.g. one file per region, client, or month).
  * --rows N    : write fixed-size chunks of N rows each.

The inverse of a merge: take one big file and break it into tidy pieces.

Usage:
    python split_excel.py <input> --by COLUMN [-o out_dir] [--format xlsx|csv]
    python split_excel.py <input> --rows N     [-o out_dir] [--format xlsx|csv]

Examples:
    python split_excel.py sales.xlsx --by region
    python split_excel.py sales.csv --rows 1000 -o chunks --format csv

Author: Synth88Labs
License: MIT
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

import pandas as pd


def read_data(path: Path, sheet: str | None) -> pd.DataFrame:
    """Read a CSV or Excel file into a DataFrame."""
    if path.suffix.lower() == ".csv":
        return pd.read_csv(path)
    sheet_arg = sheet if sheet is not None else 0
    return pd.read_excel(path, sheet_name=sheet_arg)


def safe_filename(value: object) -> str:
    """Turn a cell value into a safe file name fragment."""
    text = str(value).strip()
    if text == "" or text.lower() == "nan":
        return "blank"
    text = re.sub(r"[^\w\-]+", "_", text)  # non-word chars -> underscore
    return text.strip("_") or "blank"


def write_frame(df: pd.DataFrame, path: Path, fmt: str) -> None:
    if fmt == "csv":
        df.to_csv(path, index=False)
    else:
        df.to_excel(path, index=False, sheet_name="Sheet1")


def split_by_column(
    df: pd.DataFrame, column: str, out_dir: Path, fmt: str, prefix: str
) -> list[Path]:
    """Write one file per unique value in *column*. Returns the paths written."""
    if column not in df.columns:
        raise ValueError(f"Column '{column}' not found. Available: {', '.join(map(str, df.columns))}")

    written: list[Path] = []
    used: dict[str, int] = {}
    # sort=False keeps first-seen order; dropna=False so blanks get their own file
    for value, group in df.groupby(column, sort=False, dropna=False):
        base = safe_filename(value)
        # guard against two different values sanitizing to the same name
        if base in used:
            used[base] += 1
            base = f"{base}_{used[base]}"
        else:
            used[base] = 0
        out_path = out_dir / f"{prefix}{base}.{fmt}"
        write_frame(group, out_path, fmt)
        written.append(out_path)
    return written


def split_by_rows(
    df: pd.DataFrame, chunk: int, out_dir: Path, fmt: str, prefix: str
) -> list[Path]:
    """Write fixed-size chunks of *chunk* rows each. Returns the paths written."""
    if chunk < 1:
        raise ValueError("--rows must be 1 or greater.")
    written: list[Path] = []
    total_parts = (len(df) + chunk - 1) // chunk
    width = max(3, len(str(total_parts)))
    for i in range(total_parts):
        part = df.iloc[i * chunk : (i + 1) * chunk]
        out_path = out_dir / f"{prefix}part_{i + 1:0{width}d}.{fmt}"
        write_frame(part, out_path, fmt)
        written.append(out_path)
    return written


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Split one spreadsheet into many files (by column value or row count).",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("input", type=Path, help="Input .csv or .xlsx file.")
    mode = p.add_mutually_exclusive_group(required=True)
    mode.add_argument("--by", metavar="COLUMN", help="Split into one file per unique value in COLUMN.")
    mode.add_argument("--rows", type=int, metavar="N", help="Split into chunks of N rows each.")
    p.add_argument("-o", "--output-dir", type=Path, default=Path("split_output"),
                   help="Folder for the output files. Default: split_output")
    p.add_argument("--format", choices=["xlsx", "csv"], default="xlsx",
                   help="Output file format. Default: xlsx")
    p.add_argument("--prefix", default="", help="Optional prefix for output file names.")
    p.add_argument("--sheet", default=None, help="For Excel input: sheet name (default: first).")
    return p.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    if not args.input.is_file():
        print(f"Error: '{args.input}' is not a file.", file=sys.stderr)
        return 1

    try:
        df = read_data(args.input, args.sheet)
    except Exception as exc:  # noqa: BLE001
        print(f"Error reading '{args.input}': {exc}", file=sys.stderr)
        return 1

    if df.empty:
        print("Input has no rows — nothing to split.", file=sys.stderr)
        return 1

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

    try:
        if args.by is not None:
            written = split_by_column(df, args.by, args.output_dir, args.format, args.prefix)
            how = f"by '{args.by}'"
        else:
            written = split_by_rows(df, args.rows, args.output_dir, args.format, args.prefix)
            how = f"in chunks of {args.rows} rows"
    except ValueError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1

    print(f"Split {len(df)} rows {how} -> {len(written)} file(s) in '{args.output_dir}':")
    for path in written:
        print(f"  + {path.name}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())