ProStock/src/data/daily.py

"""Simplified daily market data interface.

A single function to fetch A股日线行情 data from Tushare.
Supports all output fields including tor (换手率) and vr (量比).
"""
import pandas as pd
from typing import Optional, List, Literal
from src.data.client import TushareClient


def get_daily(
    ts_code: str,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    trade_date: Optional[str] = None,
    adj: Literal[None, "qfq", "hfq"] = None,
    factors: Optional[List[Literal["tor", "vr"]]] = None,
    adjfactor: bool = False,
) -> pd.DataFrame:
    """Fetch daily market data for A-share stocks.

    This is a simplified interface that combines rate limiting, API calls,
    and error handling into a single function.

    Args:
        ts_code: Stock code (e.g., '000001.SZ', '600000.SH')
        start_date: Start date in YYYYMMDD format
        end_date: End date in YYYYMMDD format
        trade_date: Specific trade date in YYYYMMDD format
        adj: Adjustment type - None, 'qfq' (forward), 'hfq' (backward)
        factors: List of factors to include - 'tor' (turnover rate), 'vr' (volume ratio)
        adjfactor: Whether to include adjustment factor

    Returns:
        pd.DataFrame with daily market data containing:
        - Base fields: ts_code, trade_date, open, high, low, close, pre_close, 
                       change, pct_chg, vol, amount
        - Factor fields (if requested): tor, vr
        - Adjustment factor (if adjfactor=True): adjfactor

    Example:
        >>> data = get_daily('000001.SZ', start_date='20240101', end_date='20240131')
        >>> data = get_daily('600000.SH', factors=['tor', 'vr'])
    """
    # Initialize client
    client = TushareClient()

    # Build parameters
    params = {"ts_code": ts_code}

    if start_date:
        params["start_date"] = start_date
    if end_date:
        params["end_date"] = end_date
    if trade_date:
        params["trade_date"] = trade_date
    if adj:
        params["adj"] = adj
    if factors:
        # Tushare expects factors as comma-separated string, not list
        if isinstance(factors, list):
            factors_str = ",".join(factors)
        else:
            factors_str = factors
        params["factors"] = factors_str
    if adjfactor:
        params["adjfactor"] = "True"

    # Fetch data using pro_bar (supports factors like tor, vr)
    data = client.query("pro_bar", **params)

    return data