csvtool.py

#!/usr/bin/env python3
"""
csvtool -- Query, filter, and transform CSV files. Like csvkit, zero deps.

Chain operations on CSV data: select columns, filter rows, sort, aggregate,
convert formats. Works with any delimiter. Cross-platform.

Usage:
    py csvtool.py data.csv head 10                            # First 10 rows
    py csvtool.py data.csv cols name,age,city                 # Select columns
    py csvtool.py data.csv filter "age > 30"                  # Filter rows
    py csvtool.py data.csv sort age desc                      # Sort descending
    py csvtool.py data.csv stats                              # Column statistics
    py csvtool.py data.csv cols name,salary filter "salary > 50000" sort salary desc head 10
    py csvtool.py data.csv groupby department count           # Group + count
    py csvtool.py data.csv tojson                             # Convert to JSON
    py csvtool.py data.csv tomd                               # Convert to markdown
    py csvtool.py data.csv -d ";" cols name,age               # Semicolon delimiter
    cat data.csv | py csvtool.py head 5                       # Stdin piping

Operations (chain LEFT to RIGHT):
    head N           First N rows
    tail N           Last N rows
    cols A,B,C       Select columns (by name or 1-indexed number)
    dropcols A,B     Drop columns
    filter EXPR      Keep rows matching expression (col > val, col == val, col ~ regex)
    sort COL [desc]  Sort by column (asc default, add 'desc' for descending)
    sortnum COL [desc]  Sort numerically
    unique COL       Deduplicate by column
    count            Print row count
    stats            Column statistics (type, nulls, unique, min, max, mean)
    freq COL [N]     Top N most common values in column
    groupby COL AGG  Group by column + aggregate (count, sum:col, mean:col, min:col, max:col)
    rename OLD NEW   Rename a column
    addcol NAME EXPR Add computed column (use {col} for references)
    search TEXT      Keep rows containing text in any column
    sample N         Random sample of N rows
    tojson           Output as JSON array
    tomd             Output as markdown table
    transpose        Swap rows and columns
"""

import argparse
import csv
import io
import json
import math
import random
import re
import sys
from pathlib import Path


def read_csv(source, delimiter=",", encoding="utf-8") -> tuple[list[str], list[dict]]:
    """Read CSV into (headers, list of row dicts)."""
    if isinstance(source, str):
        with open(source, "r", encoding=encoding, errors="replace", newline="") as f:
            reader = csv.DictReader(f, delimiter=delimiter)
            headers = reader.fieldnames or []
            rows = list(reader)
    else:
        reader = csv.DictReader(source, delimiter=delimiter)
        headers = reader.fieldnames or []
        rows = list(reader)
    return list(headers), rows


def write_csv(headers: list[str], rows: list[dict], delimiter=","):
    """Write CSV to stdout."""
    writer = csv.DictWriter(sys.stdout, fieldnames=headers, delimiter=delimiter,
                            extrasaction="ignore", lineterminator="\n")
    writer.writeheader()
    for row in rows:
        writer.writerow(row)


def op_head(headers, rows, n):
    return headers, rows[:int(n)]


def op_tail(headers, rows, n):
    return headers, rows[-int(n):]


def op_cols(headers, rows, spec):
    """Select columns by name or number."""
    selected = []
    for part in spec.split(","):
        part = part.strip()
        if part.isdigit():
            idx = int(part) - 1
            if 0 <= idx < len(headers):
                selected.append(headers[idx])
        else:
            if part in headers:
                selected.append(part)
            else:
                # Fuzzy: case-insensitive match
                for h in headers:
                    if h.lower() == part.lower():
                        selected.append(h)
                        break
    return selected, [{k: r.get(k, "") for k in selected} for r in rows]


def op_dropcols(headers, rows, spec):
    drop = set(s.strip() for s in spec.split(","))
    new_headers = [h for h in headers if h not in drop]
    return new_headers, [{k: r.get(k, "") for k in new_headers} for r in rows]


def op_filter(headers, rows, expr):
    """Filter rows. Supports: col > val, col >= val, col == val, col != val, col ~ regex, col contains text."""
    # Parse expression
    m = re.match(r"(\w+)\s*(>=|<=|!=|==|>|<|~|contains)\s*(.+)", expr)
    if not m:
        print(f"Error: invalid filter expression: {expr}", file=sys.stderr)
        return headers, rows
    col, op, val = m.group(1), m.group(2), m.group(3).strip().strip("'\"")

    # Find actual column name (case-insensitive)
    actual_col = _find_col(headers, col)
    if not actual_col:
        print(f"Error: column '{col}' not found", file=sys.stderr)
        return headers, rows

    result = []
    for row in rows:
        cell = row.get(actual_col, "")
        if _compare(cell, op, val):
            result.append(row)
    return headers, result


def _compare(cell: str, op: str, val: str) -> bool:
    if op == "contains":
        return val.lower() in cell.lower()
    if op == "~":
        return bool(re.search(val, cell, re.IGNORECASE))
    if op == "==":
        return cell == val
    if op == "!=":
        return cell != val
    # Numeric comparisons
    try:
        a, b = float(cell), float(val)
    except (ValueError, TypeError):
        return cell > val if op == ">" else cell < val if op == "<" else False
    if op == ">":
        return a > b
    if op == "<":
        return a < b
    if op == ">=":
        return a >= b
    if op == "<=":
        return a <= b
    return False


def op_sort(headers, rows, col, desc=False):
    actual = _find_col(headers, col)
    if not actual:
        print(f"Error: column '{col}' not found", file=sys.stderr)
        return headers, rows
    return headers, sorted(rows, key=lambda r: r.get(actual, ""), reverse=desc)


def op_sortnum(headers, rows, col, desc=False):
    actual = _find_col(headers, col)
    if not actual:
        print(f"Error: column '{col}' not found", file=sys.stderr)
        return headers, rows

    def num_key(r):
        try:
            return float(r.get(actual, "0"))
        except (ValueError, TypeError):
            return 0.0
    return headers, sorted(rows, key=num_key, reverse=desc)


def op_unique(headers, rows, col):
    actual = _find_col(headers, col)
    if not actual:
        return headers, rows
    seen = set()
    result = []
    for r in rows:
        v = r.get(actual, "")
        if v not in seen:
            seen.add(v)
            result.append(r)
    return headers, result


def op_count(headers, rows):
    print(len(rows))
    return headers, rows


def op_stats(headers, rows):
    """Print column statistics."""
    print(f"\n  Rows: {len(rows)}  Columns: {len(headers)}\n")
    print(f"  {'Column':<20} {'Type':<8} {'Nulls':<6} {'Unique':<7} {'Min':<15} {'Max':<15} {'Mean':<10}")
    print("  " + "-" * 85)
    for h in headers:
        vals = [r.get(h, "") for r in rows]
        non_empty = [v for v in vals if v.strip()]
        nulls = len(vals) - len(non_empty)
        uniq = len(set(non_empty))

        # Detect numeric
        nums = []
        for v in non_empty:
            try:
                nums.append(float(v))
            except (ValueError, TypeError):
                pass

        if len(nums) > len(non_empty) * 0.8 and nums:
            dtype = "numeric"
            mn = f"{min(nums):g}"
            mx = f"{max(nums):g}"
            mean = f"{sum(nums)/len(nums):.2f}"
        else:
            dtype = "text"
            if non_empty:
                mn = min(non_empty)[:14]
                mx = max(non_empty)[:14]
            else:
                mn = mx = ""
            mean = ""

        print(f"  {h:<20} {dtype:<8} {nulls:<6} {uniq:<7} {mn:<15} {mx:<15} {mean:<10}")
    print()
    return headers, rows


def op_freq(headers, rows, col, n=10):
    actual = _find_col(headers, col)
    if not actual:
        return headers, rows
    counts: dict[str, int] = {}
    for r in rows:
        v = r.get(actual, "(empty)")
        counts[v] = counts.get(v, 0) + 1
    ranked = sorted(counts.items(), key=lambda x: -x[1])[:int(n)]
    total = len(rows)
    print(f"\n  Top {min(int(n), len(ranked))} values for '{actual}' ({total} rows):\n")
    for val, cnt in ranked:
        pct = cnt / total * 100 if total else 0
        bar = "#" * int(pct / 2)
        print(f"  {cnt:>6} ({pct:5.1f}%)  {bar:<25} {val}")
    print()
    return headers, rows


def op_groupby(headers, rows, col, agg_spec):
    actual = _find_col(headers, col)
    if not actual:
        return headers, rows

    groups: dict[str, list] = {}
    for r in rows:
        k = r.get(actual, "")
        groups.setdefault(k, []).append(r)

    if agg_spec == "count":
        new_headers = [actual, "count"]
        result = [{actual: k, "count": str(len(v))} for k, v in sorted(groups.items(), key=lambda x: -len(x[1]))]
    elif ":" in agg_spec:
        func, agg_col = agg_spec.split(":", 1)
        agg_actual = _find_col(headers, agg_col)
        if not agg_actual:
            print(f"Error: column '{agg_col}' not found", file=sys.stderr)
            return headers, rows
        new_headers = [actual, agg_spec]
        result = []
        for k, grp in sorted(groups.items()):
            nums = []
            for r in grp:
                try:
                    nums.append(float(r.get(agg_actual, "0")))
                except (ValueError, TypeError):
                    pass
            if func == "sum":
                val = sum(nums)
            elif func == "mean":
                val = sum(nums) / len(nums) if nums else 0
            elif func == "min":
                val = min(nums) if nums else 0
            elif func == "max":
                val = max(nums) if nums else 0
            else:
                val = len(grp)
            result.append({actual: k, agg_spec: f"{val:g}" if isinstance(val, float) else str(val)})
        result.sort(key=lambda r: float(r.get(agg_spec, "0")), reverse=True)
    else:
        return op_groupby(headers, rows, col, "count")

    return new_headers, result


def op_rename(headers, rows, old, new):
    actual = _find_col(headers, old)
    if not actual:
        return headers, rows
    new_headers = [new if h == actual else h for h in headers]
    for r in rows:
        if actual in r:
            r[new] = r.pop(actual)
    return new_headers, rows


def op_addcol(headers, rows, name, expr):
    """Add a computed column. Use {col} for references."""
    new_headers = headers + [name]
    for r in rows:
        val = expr
        for h in headers:
            val = val.replace(f"{{{h}}}", r.get(h, ""))
        # Try to evaluate as math
        try:
            val = str(eval(val, {"__builtins__": {}}, {}))
        except Exception:
            pass
        r[name] = val
    return new_headers, rows


def op_search(headers, rows, text):
    text_lower = text.lower()
    result = [r for r in rows if any(text_lower in str(v).lower() for v in r.values())]
    return headers, result


def op_sample(headers, rows, n):
    n = min(int(n), len(rows))
    return headers, random.sample(rows, n)


def op_tojson(headers, rows):
    clean = [{h: r.get(h, "") for h in headers} for r in rows]
    print(json.dumps(clean, indent=2))
    return headers, []  # Empty to suppress CSV output


def op_tomd(headers, rows):
    """Output as markdown table."""
    if not headers:
        return headers, rows
    # Calculate widths
    widths = {h: len(h) for h in headers}
    for r in rows:
        for h in headers:
            widths[h] = max(widths[h], len(str(r.get(h, ""))))

    # Header
    hline = "| " + " | ".join(h.ljust(widths[h]) for h in headers) + " |"
    sline = "| " + " | ".join("-" * widths[h] for h in headers) + " |"
    print(hline)
    print(sline)
    for r in rows:
        line = "| " + " | ".join(str(r.get(h, "")).ljust(widths[h]) for h in headers) + " |"
        print(line)
    return headers, []


def op_transpose(headers, rows):
    if not rows:
        return headers, rows
    new_headers = ["field"] + [f"row{i+1}" for i in range(len(rows))]
    new_rows = []
    for h in headers:
        row = {"field": h}
        for i, r in enumerate(rows):
            row[f"row{i+1}"] = r.get(h, "")
        new_rows.append(row)
    return new_headers, new_rows


def _find_col(headers: list[str], name: str) -> str | None:
    """Find column by exact or case-insensitive match."""
    if name in headers:
        return name
    for h in headers:
        if h.lower() == name.lower():
            return h
    if name.isdigit():
        idx = int(name) - 1
        if 0 <= idx < len(headers):
            return headers[idx]
    return None


def parse_and_run(args: list[str]):
    if not args or args[0] in ("-h", "--help"):
        print(__doc__.strip())
        return

    # Parse global flags
    delimiter = ","
    encoding = "utf-8"
    idx = 0

    while idx < len(args) and args[idx].startswith("-"):
        if args[idx] in ("-d", "--delimiter") and idx + 1 < len(args):
            delimiter = args[idx + 1]
            if delimiter == "\\t" or delimiter == "tab":
                delimiter = "\t"
            idx += 2
        elif args[idx] in ("-e", "--encoding") and idx + 1 < len(args):
            encoding = args[idx + 1]
            idx += 2
        else:
            break

    # Determine input
    input_file = None
    ops = get_op_names()

    if idx < len(args) and args[idx] not in ops:
        input_file = args[idx]
        idx += 1

    # Read input
    if input_file:
        headers, rows = read_csv(input_file, delimiter=delimiter, encoding=encoding)
    elif not sys.stdin.isatty():
        headers, rows = read_csv(io.TextIOWrapper(sys.stdin.buffer, encoding=encoding, errors="replace"),
                                 delimiter=delimiter)
    else:
        print("Error: no input. Provide a file or pipe stdin.", file=sys.stderr)
        sys.exit(1)

    # Process operations
    suppress_output = False
    while idx < len(args):
        op = args[idx].lower()
        idx += 1

        if op == "head":
            if idx < len(args) and args[idx].isdigit():
                n = int(args[idx]); idx += 1
            else:
                n = 10
            headers, rows = op_head(headers, rows, n)
        elif op == "tail":
            if idx < len(args) and args[idx].isdigit():
                n = int(args[idx]); idx += 1
            else:
                n = 10
            headers, rows = op_tail(headers, rows, n)
        elif op == "cols":
            if idx >= len(args):
                print("Error: cols requires column list", file=sys.stderr); break
            headers, rows = op_cols(headers, rows, args[idx]); idx += 1
        elif op == "dropcols":
            if idx >= len(args):
                print("Error: dropcols requires column list", file=sys.stderr); break
            headers, rows = op_dropcols(headers, rows, args[idx]); idx += 1
        elif op == "filter":
            if idx >= len(args):
                print("Error: filter requires expression", file=sys.stderr); break
            headers, rows = op_filter(headers, rows, args[idx]); idx += 1
        elif op == "sort":
            if idx >= len(args):
                print("Error: sort requires column", file=sys.stderr); break
            col = args[idx]; idx += 1
            desc = idx < len(args) and args[idx].lower() == "desc"
            if desc:
                idx += 1
            headers, rows = op_sort(headers, rows, col, desc)
        elif op == "sortnum":
            if idx >= len(args):
                print("Error: sortnum requires column", file=sys.stderr); break
            col = args[idx]; idx += 1
            desc = idx < len(args) and args[idx].lower() == "desc"
            if desc:
                idx += 1
            headers, rows = op_sortnum(headers, rows, col, desc)
        elif op == "unique":
            if idx >= len(args):
                print("Error: unique requires column", file=sys.stderr); break
            headers, rows = op_unique(headers, rows, args[idx]); idx += 1
        elif op == "count":
            op_count(headers, rows); suppress_output = True
        elif op == "stats":
            op_stats(headers, rows); suppress_output = True
        elif op == "freq":
            if idx >= len(args):
                print("Error: freq requires column", file=sys.stderr); break
            col = args[idx]; idx += 1
            n = 10
            if idx < len(args) and args[idx].isdigit():
                n = int(args[idx]); idx += 1
            op_freq(headers, rows, col, n); suppress_output = True
        elif op == "groupby":
            if idx + 1 >= len(args):
                print("Error: groupby requires COL and AGG", file=sys.stderr); break
            col = args[idx]; agg = args[idx + 1]; idx += 2
            headers, rows = op_groupby(headers, rows, col, agg)
        elif op == "rename":
            if idx + 1 >= len(args):
                print("Error: rename requires OLD NEW", file=sys.stderr); break
            headers, rows = op_rename(headers, rows, args[idx], args[idx + 1]); idx += 2
        elif op == "addcol":
            if idx + 1 >= len(args):
                print("Error: addcol requires NAME EXPR", file=sys.stderr); break
            headers, rows = op_addcol(headers, rows, args[idx], args[idx + 1]); idx += 2
        elif op == "search":
            if idx >= len(args):
                print("Error: search requires text", file=sys.stderr); break
            headers, rows = op_search(headers, rows, args[idx]); idx += 1
        elif op == "sample":
            if idx >= len(args):
                print("Error: sample requires N", file=sys.stderr); break
            headers, rows = op_sample(headers, rows, args[idx]); idx += 1
        elif op == "tojson":
            headers, rows = op_tojson(headers, rows); suppress_output = True
        elif op == "tomd":
            headers, rows = op_tomd(headers, rows); suppress_output = True
        elif op == "transpose":
            headers, rows = op_transpose(headers, rows)
        else:
            print(f"Error: unknown operation '{op}'", file=sys.stderr); break

    # Output remaining rows as CSV
    if not suppress_output and rows:
        write_csv(headers, rows, delimiter=delimiter)


def get_op_names() -> set[str]:
    return {
        "head", "tail", "cols", "dropcols", "filter", "sort", "sortnum",
        "unique", "count", "stats", "freq", "groupby", "rename", "addcol",
        "search", "sample", "tojson", "tomd", "transpose",
    }


def main():
    parse_and_run(sys.argv[1:])


if __name__ == "__main__":
    main()