yfinance-fork/cli/yfin/renderers.py

"""Output rendering: --format dispatch + --out writing.

Commands return raw Python values (DataFrame / Series / dict / list / NamedTuple).
This module turns them into bytes / text on the right sink.
"""
from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any, Optional

import pandas as pd
import yaml
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.tree import Tree

_EXT_MAP = {
    ".json": "json",
    ".ndjson": "ndjson",
    ".jsonl": "ndjson",
    ".csv": "csv",
    ".tsv": "tsv",
    ".parquet": "parquet",
    ".pq": "parquet",
    ".yaml": "yaml",
    ".yml": "yaml",
}

_VALID = {"table", "json", "ndjson", "csv", "tsv", "parquet", "yaml"}

# Module-level console — re-bound by the root app so --no-color/--quiet propagate.
_console: Optional[Console] = None


def get_console() -> Console:
    global _console
    if _console is None:
        _console = Console()
    return _console


def set_console(console: Console) -> None:
    global _console
    _console = console


def resolve_format(fmt: Optional[str], out: Optional[Path]) -> str:
    """Choose final format. Explicit --format wins; otherwise infer from --out; else default."""
    if fmt:
        if fmt not in _VALID:
            raise ValueError(f"Unknown --format {fmt!r}; valid: {sorted(_VALID)}")
        return fmt
    if out:
        return _EXT_MAP.get(out.suffix.lower(), "json")
    return "table" if sys.stdout.isatty() else "json"


def render(
    value: Any,
    fmt: Optional[str] = None,
    out: Optional[Path] = None,
    pretty: Optional[bool] = None,
    columns: Optional[str] = None,
    limit: Optional[int] = None,
) -> None:
    """Render `value` to stdout or to `out`. Mutates nothing else."""
    if value is None:
        get_console().print("[dim](no result)[/dim]" if sys.stdout.isatty() else "")
        return

    resolved = resolve_format(fmt, out)

    # Apply --columns / --limit on tabular values before rendering.
    if isinstance(value, pd.DataFrame):
        if columns:
            cols = [c.strip() for c in columns.split(",")]
            keep = [c for c in cols if c in value.columns]
            if keep:
                value = value[keep]
        if limit is not None:
            value = value.head(limit)
    elif isinstance(value, pd.Series):
        if limit is not None:
            value = value.head(limit)
    elif isinstance(value, list) and limit is not None:
        value = value[:limit]

    if out:
        out = Path(out)
        out.parent.mkdir(parents=True, exist_ok=True)

    if resolved == "table":
        _render_table(value)
    elif resolved == "json":
        _write_text(_to_json(value, pretty=pretty if pretty is not None else True), out)
    elif resolved == "ndjson":
        _write_text(_to_ndjson(value), out)
    elif resolved == "csv":
        _write_text(_to_csv(value, sep=","), out)
    elif resolved == "tsv":
        _write_text(_to_csv(value, sep="\t"), out)
    elif resolved == "parquet":
        _write_parquet(value, out)
    elif resolved == "yaml":
        _write_text(_to_yaml(value), out)
    else:
        raise ValueError(f"Unsupported format: {resolved}")


# ---------------------------------------------------------------------------- #
# Format implementations
# ---------------------------------------------------------------------------- #

def _to_json(value: Any, *, pretty: bool) -> str:
    indent = 2 if pretty else None
    if isinstance(value, pd.DataFrame):
        # orient='split' keeps index + columns + data sections, round-trippable.
        return value.to_json(orient="split", date_format="iso", indent=indent or 0)
    if isinstance(value, pd.Series):
        return value.to_json(orient="split", date_format="iso", indent=indent or 0)
    if hasattr(value, "_asdict"):  # NamedTuple
        value = {k: _coerce(v) for k, v in value._asdict().items()}
    return json.dumps(_coerce(value), default=str, indent=indent, ensure_ascii=False)


def _to_ndjson(value: Any) -> str:
    if isinstance(value, pd.DataFrame):
        return value.to_json(orient="records", lines=True, date_format="iso")
    if isinstance(value, pd.Series):
        return value.to_json(orient="records", lines=True, date_format="iso")
    if isinstance(value, list):
        return "\n".join(json.dumps(_coerce(x), default=str, ensure_ascii=False) for x in value)
    return json.dumps(_coerce(value), default=str, ensure_ascii=False)


def _to_csv(value: Any, sep: str) -> str:
    if isinstance(value, pd.DataFrame):
        return value.to_csv(sep=sep)
    if isinstance(value, pd.Series):
        return value.to_csv(sep=sep, header=True)
    if isinstance(value, list) and value and isinstance(value[0], dict):
        return pd.DataFrame(value).to_csv(sep=sep, index=False)
    raise TypeError(f"--format csv/tsv only supports DataFrames / Series / list[dict] (got {type(value).__name__})")


def _write_parquet(value: Any, out: Optional[Path]) -> None:
    if not isinstance(value, pd.DataFrame):
        if isinstance(value, pd.Series):
            value = value.to_frame()
        elif isinstance(value, list) and value and isinstance(value[0], dict):
            value = pd.DataFrame(value)
        else:
            raise TypeError("--format parquet only supports DataFrames / Series / list[dict]")
    if out is None:
        raise ValueError("--format parquet requires --out")
    value.to_parquet(out)
    get_console().print(f"[green]wrote[/green] {out}  ({len(value):,} rows × {len(value.columns)} cols)", highlight=False)


def _to_yaml(value: Any) -> str:
    if isinstance(value, pd.DataFrame):
        value = value.reset_index().to_dict(orient="records")
    elif isinstance(value, pd.Series):
        value = value.to_dict()
    return yaml.safe_dump(_coerce(value), sort_keys=False, allow_unicode=True)


def _write_text(text: str, out: Optional[Path]) -> None:
    if out:
        out.write_text(text)
        get_console().print(f"[green]wrote[/green] {out}  ({len(text):,} bytes)", highlight=False)
    else:
        # Use plain print so piping isn't corrupted by Rich markup.
        print(text)


# ---------------------------------------------------------------------------- #
# Rich table rendering
# ---------------------------------------------------------------------------- #

def _render_table(value: Any) -> None:
    c = get_console()
    if isinstance(value, pd.DataFrame):
        c.print(_dataframe_to_table(value))
    elif isinstance(value, pd.Series):
        c.print(_series_to_table(value))
    elif hasattr(value, "_asdict"):  # NamedTuple (e.g. option_chain)
        for name, sub in value._asdict().items():
            c.print(Panel(_render_any_to_renderable(sub), title=str(name)))
    elif isinstance(value, dict):
        c.print(_dict_to_table(value))
    elif isinstance(value, list):
        c.print(_list_to_renderable(value))
    else:
        c.print(value)


def _render_any_to_renderable(value: Any):
    if isinstance(value, pd.DataFrame):
        return _dataframe_to_table(value)
    if isinstance(value, pd.Series):
        return _series_to_table(value)
    if isinstance(value, dict):
        return _dict_to_table(value)
    if isinstance(value, list):
        return _list_to_renderable(value)
    return str(value)


def _dataframe_to_table(df: pd.DataFrame, max_rows: int = 200) -> Table:
    t = Table(show_header=True, header_style="bold cyan", row_styles=["", "dim"])
    index_name = df.index.name or "index"
    t.add_column(str(index_name), style="bold")
    for col in df.columns:
        t.add_column(str(col))
    n = min(len(df), max_rows)
    for idx, row in df.head(n).iterrows():
        t.add_row(_fmt(idx), *[_fmt(v) for v in row.values])
    if len(df) > n:
        t.caption = f"[dim]… {len(df) - n:,} more rows omitted[/dim]"
    return t


def _series_to_table(s: pd.Series) -> Table:
    t = Table(show_header=True, header_style="bold cyan")
    t.add_column(str(s.index.name or "index"), style="bold")
    t.add_column(str(s.name or "value"))
    for idx, val in s.items():
        t.add_row(_fmt(idx), _fmt(val))
    return t


def _dict_to_table(d: dict) -> Any:
    # If values are themselves complex, fall back to a Tree for readability.
    if any(isinstance(v, (dict, list, pd.DataFrame, pd.Series)) for v in d.values()):
        tree = Tree("[bold cyan]dict[/bold cyan]")
        _dict_to_tree(d, tree)
        return tree
    t = Table(show_header=True, header_style="bold cyan")
    t.add_column("key", style="bold")
    t.add_column("value")
    for k, v in d.items():
        t.add_row(str(k), _fmt(v))
    return t


def _dict_to_tree(d: dict, parent: Tree) -> None:
    for k, v in d.items():
        if isinstance(v, dict):
            sub = parent.add(f"[bold]{k}[/bold]")
            _dict_to_tree(v, sub)
        elif isinstance(v, pd.DataFrame):
            parent.add(f"[bold]{k}[/bold] [dim](DataFrame {v.shape[0]}×{v.shape[1]})[/dim]")
        elif isinstance(v, list):
            sub = parent.add(f"[bold]{k}[/bold] [dim](list × {len(v)})[/dim]")
            for i, item in enumerate(v[:10]):
                sub.add(_fmt(item))
            if len(v) > 10:
                sub.add(f"[dim]… {len(v) - 10} more[/dim]")
        else:
            parent.add(f"[bold]{k}[/bold]: {_fmt(v)}")


def _list_to_renderable(lst: list) -> Any:
    if not lst:
        return "[dim](empty list)[/dim]"
    if all(isinstance(x, dict) for x in lst):
        try:
            return _dataframe_to_table(pd.DataFrame(lst))
        except Exception:
            pass
    # Bullet list
    return "\n".join(f"• {_fmt(x)}" for x in lst[:200])


# ---------------------------------------------------------------------------- #
# Coercion / formatting helpers
# ---------------------------------------------------------------------------- #

def _fmt(v: Any) -> str:
    if v is None:
        return "[dim]—[/dim]"
    if isinstance(v, float):
        if v != v:  # NaN
            return "[dim]NaN[/dim]"
        if abs(v) >= 1e9:
            return f"{v:,.4g}"
        if abs(v) >= 1:
            return f"{v:,.4f}".rstrip("0").rstrip(".")
        return f"{v:.6g}"
    if isinstance(v, int):
        return f"{v:,}"
    if isinstance(v, pd.Timestamp):
        return v.isoformat()
    s = str(v)
    if len(s) > 200:
        s = s[:197] + "…"
    return s


def _coerce(value: Any) -> Any:
    """Recursively convert pandas / numpy / Timestamp into JSON-safe Python."""
    if isinstance(value, pd.DataFrame):
        return json.loads(value.to_json(orient="split", date_format="iso"))
    if isinstance(value, pd.Series):
        return json.loads(value.to_json(orient="split", date_format="iso"))
    if isinstance(value, pd.Timestamp):
        return value.isoformat()
    if isinstance(value, dict):
        return {str(k): _coerce(v) for k, v in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [_coerce(x) for x in value]
    if hasattr(value, "item"):  # numpy scalar
        try:
            return value.item()
        except Exception:
            return str(value)
    return value