dataconv.py

#!/usr/bin/env python3
"""
dataconv — Universal data format converter.

Convert between JSON, YAML, TOML, CSV, and XML with one command.
Replaces jq + yq + csvkit + xml tools with a single Python script.

Usage:
    py dataconv.py input.yaml --to json
    py dataconv.py input.json --to csv
    py dataconv.py input.toml --to yaml
    py dataconv.py data.csv --to json
    py dataconv.py config.json --to toml
    py dataconv.py input.json --to xml
    cat data.json | py dataconv.py --from json --to yaml    # stdin/stdout
    py dataconv.py input.json --to yaml -o output.yaml      # output to file
    py dataconv.py input.json --query "users[0].name"       # jq-lite query
"""

import argparse
import csv
import io
import json
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

try:
    import yaml
    HAS_YAML = True
except ImportError:
    HAS_YAML = False

# Python 3.11+ has tomllib; for writing we need tomli_w or manual
try:
    import tomllib
    HAS_TOML_READ = True
except ImportError:
    try:
        import tomli as tomllib
        HAS_TOML_READ = True
    except ImportError:
        HAS_TOML_READ = False

FORMATS = ["json", "yaml", "toml", "csv", "xml"]


# --- Readers ---

def read_json(text: str) -> any:
    return json.loads(text)


def read_yaml(text: str) -> any:
    if not HAS_YAML:
        raise ImportError("pyyaml not installed. Run: pip install pyyaml")
    return yaml.safe_load(text)


def read_toml(text: str) -> any:
    if not HAS_TOML_READ:
        raise ImportError("tomllib not available. Use Python 3.11+ or: pip install tomli")
    return tomllib.loads(text)


def read_csv_data(text: str) -> list[dict]:
    reader = csv.DictReader(io.StringIO(text))
    return list(reader)


def read_xml(text: str) -> dict:
    """Parse XML to dict (simple recursive conversion)."""
    root = ET.fromstring(text)
    return {root.tag: _xml_element_to_dict(root)}


def _xml_element_to_dict(elem) -> dict | str:
    result = {}
    if elem.attrib:
        result["@attributes"] = dict(elem.attrib)
    children = list(elem)
    if not children:
        text = (elem.text or "").strip()
        if result:
            if text:
                result["#text"] = text
            return result
        return text
    for child in children:
        child_data = _xml_element_to_dict(child)
        if child.tag in result:
            if not isinstance(result[child.tag], list):
                result[child.tag] = [result[child.tag]]
            result[child.tag].append(child_data)
        else:
            result[child.tag] = child_data
    return result


# --- Writers ---

def write_json(data: any, indent: int = 2) -> str:
    return json.dumps(data, indent=indent, ensure_ascii=False, default=str)


def write_yaml(data: any) -> str:
    if not HAS_YAML:
        raise ImportError("pyyaml not installed. Run: pip install pyyaml")
    return yaml.dump(data, default_flow_style=False, allow_unicode=True, sort_keys=False)


def write_toml(data: dict) -> str:
    """Simple TOML writer (handles common cases)."""
    lines = []
    _write_toml_section(data, lines, [])
    return "\n".join(lines) + "\n"


def _write_toml_section(data: dict, lines: list, path: list):
    # Write simple key-values first, then tables
    simple = {}
    tables = {}
    for k, v in data.items():
        if isinstance(v, dict):
            tables[k] = v
        elif isinstance(v, list) and v and isinstance(v[0], dict):
            tables[k] = v
        else:
            simple[k] = v

    for k, v in simple.items():
        lines.append(f"{k} = {_toml_value(v)}")

    for k, v in tables.items():
        if isinstance(v, list):
            for item in v:
                lines.append(f"\n[[{'.'.join(path + [k])}]]")
                _write_toml_section(item, lines, path + [k])
        else:
            lines.append(f"\n[{'.'.join(path + [k])}]")
            _write_toml_section(v, lines, path + [k])


def _toml_value(v) -> str:
    if isinstance(v, str):
        return f'"{v}"'
    elif isinstance(v, bool):
        return "true" if v else "false"
    elif isinstance(v, (int, float)):
        return str(v)
    elif isinstance(v, list):
        return "[" + ", ".join(_toml_value(i) for i in v) + "]"
    else:
        return f'"{v}"'


def write_csv(data: any) -> str:
    """Convert data to CSV. Expects list of dicts or list of lists."""
    if not isinstance(data, list):
        if isinstance(data, dict):
            # Single dict: convert to list of one-row
            data = [data]
        else:
            raise ValueError("CSV output requires a list of records")

    if not data:
        return ""

    output = io.StringIO()
    if isinstance(data[0], dict):
        fieldnames = list(data[0].keys())
        writer = csv.DictWriter(output, fieldnames=fieldnames)
        writer.writeheader()
        for row in data:
            writer.writerow({k: str(v) for k, v in row.items()})
    else:
        writer = csv.writer(output)
        for row in data:
            writer.writerow(row)

    return output.getvalue()


def write_xml(data: any, root_tag: str = "root") -> str:
    """Convert dict/list to XML string."""
    if isinstance(data, dict):
        if len(data) == 1:
            root_tag = list(data.keys())[0]
            data = data[root_tag]

    root = _dict_to_xml_element(root_tag, data)
    ET.indent(root, space="  ")
    return ET.tostring(root, encoding="unicode", xml_declaration=True) + "\n"


def _dict_to_xml_element(tag: str, data) -> ET.Element:
    elem = ET.Element(tag)
    if isinstance(data, dict):
        for k, v in data.items():
            if k == "@attributes":
                for ak, av in v.items():
                    elem.set(ak, str(av))
            elif k == "#text":
                elem.text = str(v)
            elif isinstance(v, list):
                for item in v:
                    child = _dict_to_xml_element(k, item)
                    elem.append(child)
            else:
                child = _dict_to_xml_element(k, v)
                elem.append(child)
    elif isinstance(data, list):
        for item in data:
            child = _dict_to_xml_element("item", item)
            elem.append(child)
    else:
        elem.text = str(data)
    return elem


# --- Query (jq-lite) ---

def query_data(data: any, query: str) -> any:
    """Simple dot-notation query: 'users[0].name' or 'config.database.host'"""
    parts = []
    for part in query.replace("]", "").split("."):
        if "[" in part:
            key, idx = part.split("[")
            if key:
                parts.append(key)
            parts.append(int(idx))
        else:
            parts.append(part)

    result = data
    for part in parts:
        if isinstance(part, int):
            result = result[part]
        elif isinstance(result, dict):
            result = result[part]
        else:
            raise KeyError(f"Cannot access '{part}' on {type(result).__name__}")
    return result


# --- Format detection ---

def detect_format(filepath: str) -> str | None:
    ext = Path(filepath).suffix.lower()
    mapping = {
        ".json": "json", ".yaml": "yaml", ".yml": "yaml",
        ".toml": "toml", ".csv": "csv", ".xml": "xml",
    }
    return mapping.get(ext)


# --- Main ---

def main():
    parser = argparse.ArgumentParser(
        description="dataconv -- universal data format converter (JSON/YAML/TOML/CSV/XML)"
    )
    parser.add_argument("input", nargs="?", help="Input file (or - for stdin)")
    parser.add_argument("--to", "-t", choices=FORMATS, help="Output format")
    parser.add_argument("--from", "-f", dest="from_fmt", choices=FORMATS, help="Input format (auto-detected from extension)")
    parser.add_argument("--output", "-o", help="Output file (default: stdout)")
    parser.add_argument("--query", "-q", help="Query expression (dot notation, e.g. 'users[0].name')")
    parser.add_argument("--compact", "-c", action="store_true", help="Compact JSON output (no indentation)")

    args = parser.parse_args()

    # Read input
    if args.input and args.input != "-":
        filepath = Path(args.input)
        if not filepath.exists():
            print(f"Error: {filepath} not found", file=sys.stderr)
            sys.exit(1)
        text = filepath.read_text(encoding="utf-8")
        in_fmt = args.from_fmt or detect_format(str(filepath))
    else:
        text = sys.stdin.read()
        in_fmt = args.from_fmt
        if not in_fmt:
            print("Error: specify --from format when reading from stdin", file=sys.stderr)
            sys.exit(1)

    if not in_fmt:
        print("Error: cannot detect input format. Use --from to specify.", file=sys.stderr)
        sys.exit(1)

    # Parse
    readers = {"json": read_json, "yaml": read_yaml, "toml": read_toml, "csv": read_csv_data, "xml": read_xml}
    try:
        data = readers[in_fmt](text)
    except Exception as e:
        print(f"Error parsing {in_fmt}: {e}", file=sys.stderr)
        sys.exit(1)

    # Query
    if args.query:
        try:
            data = query_data(data, args.query)
        except (KeyError, IndexError, TypeError) as e:
            print(f"Query error: {e}", file=sys.stderr)
            sys.exit(1)

    # Determine output format
    out_fmt = args.to
    if not out_fmt:
        if args.output:
            out_fmt = detect_format(args.output)
        if not out_fmt:
            out_fmt = "json"  # default

    # Write
    writers = {
        "json": lambda d: write_json(d, indent=0 if args.compact else 2),
        "yaml": write_yaml,
        "toml": write_toml,
        "csv": write_csv,
        "xml": write_xml,
    }
    try:
        output = writers[out_fmt](data)
    except Exception as e:
        print(f"Error converting to {out_fmt}: {e}", file=sys.stderr)
        sys.exit(1)

    # Output
    if args.output:
        Path(args.output).write_text(output, encoding="utf-8")
        print(f"Written to {args.output}", file=sys.stderr)
    else:
        print(output, end="")


if __name__ == "__main__":
    main()