tabular_convert.py

#!/usr/bin/env python3
"""Convert between CSV and JSON without losing structure.

Two directions:
    csv2json  rows -> list of objects, with optional type inference and
              dotted-key expansion into nested objects
    json2csv  list of objects -> rows, flattening nested objects/arrays
              into dotted keys so nothing is silently dropped

Round-tripping is the design goal: json2csv followed by csv2json --nest
--infer should give you back what you started with for any JSON that is a
list of flat-or-nested objects with scalar leaves.

Usage:
    tabular_convert.py csv2json input.csv -o out.json --infer --nest
    tabular_convert.py json2csv input.json -o out.csv
    cat in.csv | tabular_convert.py csv2json - -o -
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from typing import Any, Iterable, Iterator

# csv.field_size_limit(sys.maxsize) overflows on some platforms; step down
# until the C long accepts it.
def _raise_field_limit() -> None:
    limit = sys.maxsize
    while True:
        try:
            csv.field_size_limit(limit)
            return
        except OverflowError:
            limit //= 2


_raise_field_limit()

# CSV has no type system, so a few values are encoded as literal text on the
# way out and decoded on the way back in. An empty cell is reserved to mean
# "this record has no such field", which is what lets ragged records survive.
SENTINEL_ABSENT = ""
SENTINEL_NULL = "null"
SENTINEL_EMPTY_OBJ = "{}"
SENTINEL_EMPTY_LIST = "[]"


class ConversionError(Exception):
    """Raised for malformed input that the user needs to fix."""


# --------------------------------------------------------------------------
# type inference
# --------------------------------------------------------------------------

_TRUE = {"true", "True", "TRUE"}
_FALSE = {"false", "False", "FALSE"}
_NULL = {"null", "NULL", "None"}


def infer_scalar(text: str) -> Any:
    """Turn a CSV cell into int/float/bool/None when it unambiguously is one.

    Deliberately conservative. A leading zero ("007"), a leading "+", or
    surrounding whitespace means the author probably wanted a string, so it
    stays a string. This keeps zip codes and phone numbers intact.
    """
    if text == "":
        return None
    if text in _TRUE:
        return True
    if text in _FALSE:
        return False
    if text in _NULL:
        return None
    if text != text.strip():
        return text
    # int()/float() accept "1_000" and "+5"; JSON does not, and a user who
    # typed either almost certainly meant a literal string.
    if "_" in text or text.startswith("+"):
        return text

    body = text[1:] if text.startswith("-") else text
    if not body:
        return text
    if body[0] == "0" and len(body) > 1 and body[1] != ".":
        return text  # 007, 0123 -> keep as string

    if body.isdigit():
        return int(text)
    try:
        value = float(text)
    except ValueError:
        return text
    # float() accepts "nan"/"inf", which are not valid JSON. Keep them strings.
    if value != value or value in (float("inf"), float("-inf")):
        return text
    return value


# --------------------------------------------------------------------------
# nesting / flattening
# --------------------------------------------------------------------------


def encode_cell(value: Any) -> str:
    """Render a JSON scalar as CSV text that infer_scalar can read back.

    Python's str() would emit True/False/None, which no other CSV consumer
    recognises, so JSON literals are used instead.
    """
    if value is None:
        return SENTINEL_NULL
    if value is True:
        return "true"
    if value is False:
        return "false"
    return str(value)


def flatten(obj: Any, prefix: str = "") -> dict[str, Any]:
    """Flatten nested dicts/lists into dotted keys.

    {"a": {"b": 1}, "c": [10, 20]} -> {"a.b": 1, "c.0": 10, "c.1": 20}

    Empty dicts and empty lists have no leaves, so they would vanish. They
    are emitted as their JSON literal to survive the round trip.
    """
    flat: dict[str, Any] = {}
    if isinstance(obj, dict):
        if not obj and prefix:
            flat[prefix] = SENTINEL_EMPTY_OBJ
            return flat
        for key, value in obj.items():
            path = f"{prefix}.{key}" if prefix else str(key)
            flat.update(flatten(value, path))
    elif isinstance(obj, list):
        if not obj and prefix:
            flat[prefix] = SENTINEL_EMPTY_LIST
            return flat
        for index, value in enumerate(obj):
            path = f"{prefix}.{index}" if prefix else str(index)
            flat.update(flatten(value, path))
    else:
        flat[prefix] = obj
    return flat


def _is_index(key: str) -> bool:
    return key.isdigit()


def nest(flat: dict[str, Any]) -> dict[str, Any]:
    """Inverse of flatten. Dotted keys become nested dicts.

    Numeric path segments rebuild lists, so "c.0"/"c.1" comes back as a list
    rather than a dict with "0"/"1" keys.

    A ragged CSV can carry both "x" and "x.0" as columns, because some rows
    had a populated list there and others had an empty one. The populated
    branch wins; the blank scalar at "x" is padding from DictWriter's
    restval, not real data.
    """
    root: dict[str, Any] = {}
    for dotted in flat:
        value = flat[dotted]
        # An empty cell means this record simply has no such field. Ragged
        # records share one union header, so most rows have empty cells for
        # columns belonging to other rows; materialising those as nulls would
        # invent keys the source never had. An explicit JSON null arrives as
        # the "null" sentinel instead, so the two stay distinguishable.
        if value == SENTINEL_ABSENT:
            continue

        parts = dotted.split(".")
        cursor: dict[str, Any] = root
        for depth, part in enumerate(parts[:-1]):
            nxt = cursor.get(part)
            if not isinstance(nxt, dict):
                if nxt is not None:
                    raise ConversionError(
                        f"column {dotted!r} conflicts with column "
                        f"{'.'.join(parts[:depth + 1])!r}, which holds a value"
                    )
                nxt = {}
                cursor[part] = nxt
            cursor = nxt

        leaf = parts[-1]
        if value == SENTINEL_EMPTY_OBJ:
            value = {}
        elif value == SENTINEL_EMPTY_LIST:
            value = []
        cursor[leaf] = value
    return _listify(root)


def _listify(node: Any) -> Any:
    """Convert dicts whose keys are exactly 0..n-1 back into lists."""
    if not isinstance(node, dict):
        return node
    converted = {k: _listify(v) for k, v in node.items()}
    if converted and all(_is_index(k) for k in converted):
        indices = sorted(int(k) for k in converted)
        if indices == list(range(len(indices))):
            return [converted[str(i)] for i in indices]
    return converted


# --------------------------------------------------------------------------
# conversion
# --------------------------------------------------------------------------


def csv_to_records(
    stream: Iterable[str], *, do_infer: bool = False, do_nest: bool = False,
    delimiter: str | None = None, string_columns: Iterable[str] = (),
) -> Iterator[dict[str, Any]]:
    """Stream CSV rows as dicts. Yields lazily so large files stay cheap.

    Columns named in string_columns are exempt from type inference, which is
    how you keep a zip code like 90210 from becoming an integer.
    """
    keep_as_string = set(string_columns)
    if delimiter is None:
        reader = csv.DictReader(stream)
    else:
        reader = csv.DictReader(stream, delimiter=delimiter)

    if reader.fieldnames is None:
        return
    if any(name is None or name == "" for name in reader.fieldnames):
        raise ConversionError("CSV header has an empty column name")
    duplicates = {n for n in reader.fieldnames if reader.fieldnames.count(n) > 1}
    if duplicates:
        raise ConversionError(
            f"CSV header has duplicate column names: {sorted(duplicates)}"
        )
    # A typo here would silently fail to protect the column, so say so.
    unknown = keep_as_string - set(reader.fieldnames)
    if unknown:
        raise ConversionError(
            f"--string-columns names columns not in the header: "
            f"{sorted(unknown)}"
        )

    for row in reader:
        # A short row leaves DictReader values as None; a long row dumps the
        # overflow under the None key. Both mean the row does not match the
        # header, which is worth surfacing rather than silently coercing.
        if None in row:
            raise ConversionError(
                f"row has more fields than the header: {row[None]!r}"
            )
        clean = {k: ("" if v is None else v) for k, v in row.items()}
        if do_nest:
            # Drop absent fields before inference, otherwise "" and the
            # "null" sentinel both become None and the difference between
            # "field missing" and "field is null" is lost.
            clean = {k: v for k, v in clean.items() if v != SENTINEL_ABSENT}
        if do_infer:
            clean = {
                k: v if k in keep_as_string else infer_scalar(v)
                for k, v in clean.items()
            }
        yield nest(clean) if do_nest else clean


def records_to_csv(records: list[Any], out, *, delimiter: str = ",") -> int:
    """Write records as CSV. Returns the row count.

    The header is the union of every record's keys, in first-seen order, so
    ragged records do not lose columns.
    """
    if not isinstance(records, list):
        raise ConversionError(
            "json2csv needs a JSON array of objects at the top level, "
            f"got {type(records).__name__}"
        )
    flat_rows: list[dict[str, Any]] = []
    header: list[str] = []
    seen: set[str] = set()
    for i, record in enumerate(records):
        if not isinstance(record, dict):
            raise ConversionError(
                f"element {i} is {type(record).__name__}, expected an object"
            )
        flat = flatten(record)
        flat_rows.append(flat)
        for key in flat:
            if key not in seen:
                seen.add(key)
                header.append(key)

    writer = csv.DictWriter(
        out, fieldnames=header, delimiter=delimiter,
        restval=SENTINEL_ABSENT, lineterminator="\n",
    )
    writer.writeheader()
    for flat in flat_rows:
        writer.writerow({k: encode_cell(v) for k, v in flat.items()})
    return len(flat_rows)


# --------------------------------------------------------------------------
# cli
# --------------------------------------------------------------------------


def _open_in(path: str):
    if path == "-":
        return sys.stdin
    return open(path, "r", newline="", encoding="utf-8")


def _open_out(path: str):
    if path == "-":
        return sys.stdout
    return open(path, "w", newline="", encoding="utf-8")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="tabular_convert.py",
        description="Convert between CSV and JSON. Use - for stdin/stdout.",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    c2j = sub.add_parser("csv2json", help="CSV -> JSON array of objects")
    c2j.add_argument("input")
    c2j.add_argument("-o", "--output", default="-")
    c2j.add_argument("--infer", action="store_true",
                     help="convert numeric/boolean/empty cells to real types")
    c2j.add_argument("--nest", action="store_true",
                     help="expand dotted column names into nested objects")
    c2j.add_argument("--delimiter", default=None,
                     help="input delimiter (default: ,)")
    c2j.add_argument("--string-columns", default="",
                     help="comma-separated columns to exempt from --infer, "
                          "e.g. zip,phone,account_id")
    c2j.add_argument("--indent", type=int, default=2,
                     help="output indent, 0 for compact")

    j2c = sub.add_parser("json2csv", help="JSON array of objects -> CSV")
    j2c.add_argument("input")
    j2c.add_argument("-o", "--output", default="-")
    j2c.add_argument("--delimiter", default=",")

    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        if args.command == "csv2json":
            src = _open_in(args.input)
            try:
                records = list(csv_to_records(
                    src, do_infer=args.infer, do_nest=args.nest,
                    delimiter=args.delimiter,
                    string_columns=[
                        c for c in args.string_columns.split(",") if c
                    ],
                ))
            finally:
                if src is not sys.stdin:
                    src.close()
            dst = _open_out(args.output)
            try:
                json.dump(records, dst,
                          indent=args.indent if args.indent > 0 else None,
                          ensure_ascii=False)
                dst.write("\n")
            finally:
                if dst is not sys.stdout:
                    dst.close()
            print(f"wrote {len(records)} records", file=sys.stderr)

        else:
            src = _open_in(args.input)
            try:
                payload = json.load(src)
            except json.JSONDecodeError as exc:
                raise ConversionError(f"input is not valid JSON: {exc}") from exc
            finally:
                if src is not sys.stdin:
                    src.close()
            dst = _open_out(args.output)
            try:
                count = records_to_csv(payload, dst, delimiter=args.delimiter)
            finally:
                if dst is not sys.stdout:
                    dst.close()
            print(f"wrote {count} rows", file=sys.stderr)

    except ConversionError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    except FileNotFoundError as exc:
        print(f"error: no such file: {exc.filename}", file=sys.stderr)
        return 2
    return 0


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