yfinance-fork/cli/yfin/commands/screen.py

"""`yfin screen ...` — predefined + custom market screens."""
from __future__ import annotations

import json
import sys
from enum import Enum
from pathlib import Path
from typing import List, Optional

import typer
import yfinance as yf

from .._shared import render_via_state


class QueryType(str, Enum):
    equity = "equity"; mutualfund = "mutualfund"; etf = "etf"


_QUERY_CLASS = {
    QueryType.equity: yf.EquityQuery,
    QueryType.mutualfund: yf.FundQuery,
    QueryType.etf: yf.ETFQuery,
}


app = typer.Typer(no_args_is_help=True, help="Predefined + custom market screens.")


@app.command("list-predefined")
def list_predefined(
    ctx: typer.Context,
    include_bodies: bool = typer.Option(False, "--include-bodies/--no-include-bodies"),
):
    """Show available predefined Yahoo screener names."""
    if include_bodies:
        out = {k: {"sortField": v.get("sortField"), "sortType": v.get("sortType"),
                   "query": str(v["query"])} for k, v in yf.PREDEFINED_SCREENER_QUERIES.items()}
    else:
        out = sorted(yf.PREDEFINED_SCREENER_QUERIES.keys())
    render_via_state(ctx, out)


@app.command("predefined")
def predefined(
    ctx: typer.Context,
    name: str = typer.Argument(..., help="Predefined screener key (e.g. day_gainers)."),
    count: int = typer.Option(25, "--count"),
    offset: Optional[int] = typer.Option(None, "--offset"),
    sort_field: Optional[str] = typer.Option(None, "--sort-field"),
    sort_asc: Optional[bool] = typer.Option(None, "--sort-asc/--sort-desc"),
):
    """Run a predefined Yahoo screener."""
    if name not in yf.PREDEFINED_SCREENER_QUERIES:
        raise typer.BadParameter(
            f"Unknown predefined screener {name!r}. Use `yfin screen list-predefined`."
        )
    out = yf.screen(
        name, count=count, offset=offset,
        sortField=sort_field, sortAsc=sort_asc,
    )
    render_via_state(ctx, out)


@app.command("run")
def run(
    ctx: typer.Context,
    query: str = typer.Option(
        ..., "--query", "-Q",
        help='JSON QueryBase tree, "@file.json", or "-" to read stdin.',
    ),
    type: QueryType = typer.Option(QueryType.equity, "--type"),
    size: int = typer.Option(100, "--size"),
    offset: int = typer.Option(0, "--offset"),
    sort_field: Optional[str] = typer.Option(None, "--sort-field"),
    sort_asc: bool = typer.Option(False, "--sort-asc/--sort-desc"),
):
    """Run a custom screen built from a QueryBase tree."""
    raw = _load_query(query)
    q = _build_query(type, raw)
    out = yf.screen(
        q, size=size, offset=offset,
        sortField=sort_field, sortAsc=sort_asc,
    )
    render_via_state(ctx, out)


@app.command("build")
def build(
    ctx: typer.Context,
    type: QueryType = typer.Option(..., "--type"),
    operator: str = typer.Option(..., "--operator"),
    operands: str = typer.Option(..., "--operands", help='JSON array. For and/or, items are nested {operator, operands} dicts; for value ops, items are field+scalar.'),
):
    """Build and emit a QueryBase serialization (no Yahoo call)."""
    cls = _QUERY_CLASS[type]
    ops = json.loads(operands)
    if not isinstance(ops, list):
        raise typer.BadParameter("--operands must be a JSON array.")
    op_up = operator.upper()
    if op_up in {"AND", "OR"}:
        wrapped = [_build_query(type, item) if isinstance(item, dict) else item for item in ops]
        q = cls(operator, wrapped)
    else:
        q = cls(operator, ops)
    render_via_state(ctx, {"type": type.value, "query": q.to_dict()})


def _load_query(query: str) -> dict:
    if query == "-":
        text = sys.stdin.read()
    elif query.startswith("@"):
        text = Path(query[1:]).read_text()
    else:
        text = query
    raw = json.loads(text)
    if "query" in raw and isinstance(raw["query"], dict):
        raw = raw["query"]
    return raw


def _build_query(type: QueryType, raw: dict):
    cls = _QUERY_CLASS[type]
    op = raw["operator"]
    operands = raw.get("operands") or raw.get("operand") or []
    built: List = []
    for o in operands:
        if isinstance(o, dict) and "operator" in o:
            built.append(_build_query(type, o))
        else:
            built.append(o)
    return cls(op, built)