StockMind/stockmind/core/market_data.py

from __future__ import annotations
import time
from datetime import datetime, timedelta
from typing import Optional
import pandas as pd
import yfinance as yf

_quote_cache: dict[str, tuple[dict, float]] = {}
_CACHE_TTL = 60  # seconds


def get_quote(symbol: str, force: bool = False) -> dict:
    now = time.time()
    if not force and symbol in _quote_cache:
        data, ts = _quote_cache[symbol]
        if now - ts < _CACHE_TTL:
            return data

    try:
        ticker = yf.Ticker(symbol)
        info = ticker.fast_info
        hist = ticker.history(period="2d", interval="1d")

        price = float(info.last_price or 0)
        prev_close = float(info.previous_close or price)
        change = price - prev_close
        change_pct = (change / prev_close * 100) if prev_close else 0.0

        data = {
            "symbol": symbol,
            "price": price,
            "change": change,
            "change_pct": change_pct,
            "volume": int(info.three_month_average_volume or 0),
            "market_cap": getattr(info, "market_cap", None),
            "prev_close": prev_close,
            "day_high": float(getattr(info, "day_high", price) or price),
            "day_low": float(getattr(info, "day_low", price) or price),
            "fifty_two_week_high": float(getattr(info, "fifty_two_week_high", 0) or 0),
            "fifty_two_week_low": float(getattr(info, "fifty_two_week_low", 0) or 0),
            "error": None,
        }
    except Exception as e:
        data = {
            "symbol": symbol, "price": 0.0, "change": 0.0, "change_pct": 0.0,
            "volume": 0, "market_cap": None, "prev_close": 0.0,
            "day_high": 0.0, "day_low": 0.0,
            "fifty_two_week_high": 0.0, "fifty_two_week_low": 0.0,
            "error": str(e),
        }

    _quote_cache[symbol] = (data, now)
    return data


def get_history(symbol: str, period: str = "6mo", interval: str = "1d") -> pd.DataFrame:
    try:
        ticker = yf.Ticker(symbol)
        df = ticker.history(period=period, interval=interval)
        df.index = pd.to_datetime(df.index)
        return df
    except Exception:
        return pd.DataFrame()


def get_fundamentals(symbol: str) -> dict:
    try:
        ticker = yf.Ticker(symbol)
        info = ticker.info
        return {
            "name": info.get("longName", symbol),
            "sector": info.get("sector", ""),
            "industry": info.get("industry", ""),
            "market_cap": info.get("marketCap"),
            "pe_ratio": info.get("trailingPE"),
            "forward_pe": info.get("forwardPE"),
            "eps": info.get("trailingEps"),
            "revenue": info.get("totalRevenue"),
            "profit_margin": info.get("profitMargins"),
            "dividend_yield": info.get("dividendYield"),
            "beta": info.get("beta"),
            "week_52_high": info.get("fiftyTwoWeekHigh"),
            "week_52_low": info.get("fiftyTwoWeekLow"),
            "avg_volume": info.get("averageVolume"),
            "description": info.get("longBusinessSummary", ""),
        }
    except Exception:
        return {"name": symbol, "sector": "", "industry": "", "error": True}


def search_symbols(query: str) -> list[dict]:
    if not query or len(query) < 1:
        return []
    try:
        results = yf.Search(query, max_results=10)
        quotes = results.quotes if hasattr(results, "quotes") else []
        return [
            {"symbol": q.get("symbol", ""), "name": q.get("shortname", q.get("longname", ""))}
            for q in quotes
            if q.get("symbol")
        ]
    except Exception:
        return []


PERIOD_MAP = {
    "1D": ("1d", "5m"),
    "1W": ("5d", "15m"),
    "1M": ("1mo", "1h"),
    "3M": ("3mo", "1d"),
    "6M": ("6mo", "1d"),
    "1Y": ("1y", "1d"),
    "5Y": ("5y", "1wk"),
}


def get_chart_data(symbol: str, period_label: str = "6M") -> pd.DataFrame:
    period, interval = PERIOD_MAP.get(period_label, ("6mo", "1d"))
    return get_history(symbol, period=period, interval=interval)


def get_batch_quotes(symbols: list[str]) -> dict[str, dict]:
    results = {}
    for sym in symbols:
        results[sym] = get_quote(sym)
    return results


SCREENER_UNIVERSE = [
    "AAPL", "MSFT", "GOOGL", "AMZN", "NVDA", "META", "TSLA", "BRK-B", "UNH", "JPM",
    "V", "XOM", "JNJ", "PG", "MA", "HD", "CVX", "MRK", "ABBV", "PEP",
    "KO", "AVGO", "COST", "LLY", "MCD", "TMO", "ACN", "BAC", "CSCO", "WMT",
    "ABT", "CRM", "DIS", "NFLX", "AMD", "INTC", "QCOM", "TXN", "PYPL", "AMGN",
    "SPY", "QQQ", "DIA", "GLD", "SLV",
    "BTC-USD", "ETH-USD", "SOL-USD", "BNB-USD", "ADA-USD",
]


def get_screener_data(symbols: list[str] | None = None) -> list[dict]:
    if symbols is None:
        symbols = SCREENER_UNIVERSE
    results = []
    for sym in symbols:
        try:
            ticker = yf.Ticker(sym)
            info = ticker.info
            fast = ticker.fast_info
            hist = ticker.history(period="1y", interval="1d")
            rsi_val = None
            if len(hist) >= 14:
                delta = hist["Close"].diff()
                gain = delta.clip(lower=0).rolling(14).mean()
                loss = (-delta.clip(upper=0)).rolling(14).mean()
                rs = gain / loss
                rsi_series = 100 - (100 / (1 + rs))
                rsi_val = round(float(rsi_series.iloc[-1]), 1) if not rsi_series.empty else None

            price = float(fast.last_price or 0)
            week_52_high = float(getattr(fast, "year_high", 0) or 0)
            week_52_low = float(getattr(fast, "year_low", 0) or 0)

            results.append({
                "symbol": sym,
                "name": info.get("shortName", sym),
                "price": price,
                "sector": info.get("sector", ""),
                "market_cap": info.get("marketCap"),
                "pe_ratio": info.get("trailingPE"),
                "rsi": rsi_val,
                "week_52_high": week_52_high,
                "week_52_low": week_52_low,
                "pct_from_52h": round((price - week_52_high) / week_52_high * 100, 1) if week_52_high else None,
                "pct_from_52l": round((price - week_52_low) / week_52_low * 100, 1) if week_52_low else None,
            })
        except Exception:
            continue
    return results