77 lines
1.7 KiB
Python
77 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
|
|
import pandas as pd
|
|
from dotenv import load_dotenv
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TushareConfig:
|
|
token: str
|
|
timeout: int = 30
|
|
|
|
|
|
def load_tushare_config(env_path: str | None = None) -> TushareConfig:
|
|
if env_path:
|
|
load_dotenv(env_path)
|
|
else:
|
|
load_dotenv()
|
|
|
|
token = os.getenv("TUSHARE_TOKEN", "").strip()
|
|
if not token:
|
|
raise RuntimeError("TUSHARE_TOKEN is required (set it in .env)")
|
|
|
|
timeout_s = os.getenv("TUSHARE_TIMEOUT", "30").strip() or "30"
|
|
try:
|
|
timeout = int(timeout_s)
|
|
except ValueError:
|
|
timeout = 30
|
|
|
|
return TushareConfig(token=token, timeout=timeout)
|
|
|
|
|
|
def pro_api(cfg: TushareConfig):
|
|
import tushare as ts
|
|
|
|
ts.set_token(cfg.token)
|
|
return ts.pro_api(timeout=cfg.timeout)
|
|
|
|
|
|
def fetch_stock_daily(
|
|
cfg: TushareConfig,
|
|
ts_code: str | None = None,
|
|
trade_date: str | None = None,
|
|
start_date: str | None = None,
|
|
end_date: str | None = None,
|
|
fields: str | None = None,
|
|
) -> pd.DataFrame:
|
|
api = pro_api(cfg)
|
|
return api.daily(
|
|
ts_code=ts_code,
|
|
trade_date=trade_date,
|
|
start_date=start_date,
|
|
end_date=end_date,
|
|
fields=fields,
|
|
)
|
|
|
|
|
|
def fetch_fund_daily(
|
|
cfg: TushareConfig,
|
|
ts_code: str | None = None,
|
|
trade_date: str | None = None,
|
|
start_date: str | None = None,
|
|
end_date: str | None = None,
|
|
fields: str | None = None,
|
|
) -> pd.DataFrame:
|
|
"""Fetch ETF/fund daily bars via Tushare Pro `fund_daily`."""
|
|
api = pro_api(cfg)
|
|
return api.fund_daily(
|
|
ts_code=ts_code,
|
|
trade_date=trade_date,
|
|
start_date=start_date,
|
|
end_date=end_date,
|
|
fields=fields,
|
|
)
|