67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Download OHLCV data from yfinance and save as CSV for the C++ engine.
|
|
|
|
Usage:
|
|
python scripts/download_data.py AAPL MSFT --start 2015-01-01 --end 2024-01-01
|
|
python scripts/download_data.py AAPL,MSFT --start 2015-01-01 --end 2024-01-01
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import yfinance as yf
|
|
import pandas as pd
|
|
|
|
DATA_DIR = os.path.join(
|
|
os.environ.get("XDG_DATA_HOME", os.path.expanduser("~/.local/share")),
|
|
"backtest",
|
|
"data",
|
|
)
|
|
|
|
|
|
def download(symbols: list[str], start: str, end: str, out_dir: str = "data"):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
# download all at once (one request)
|
|
raw = yf.download(symbols, start=start, end=end, auto_adjust=True, progress=False)
|
|
|
|
for sym in symbols:
|
|
print(f"Downloading {sym}...")
|
|
|
|
if len(symbols) == 1:
|
|
df = raw[["Open", "High", "Low", "Close", "Volume"]].copy()
|
|
else:
|
|
# multi-symbol: yfinance returns MultiIndex columns (field, symbol)
|
|
df = raw.xs(sym, axis=1, level=1)[
|
|
["Open", "High", "Low", "Close", "Volume"]
|
|
].copy()
|
|
|
|
if df.empty:
|
|
print(f" WARNING: no data for {sym}")
|
|
continue
|
|
|
|
df.index.name = "date"
|
|
df.columns = ["open", "high", "low", "close", "volume"]
|
|
df["volume"] = df["volume"].astype(int)
|
|
|
|
path = os.path.join(out_dir, f"{sym}.csv")
|
|
df.to_csv(path)
|
|
print(f" Saved {len(df)} bars -> {path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("symbols", nargs="+")
|
|
parser.add_argument("--start", default="2015-01-01")
|
|
parser.add_argument("--end", default="2024-01-01")
|
|
parser.add_argument("--outdir", default=DATA_DIR)
|
|
args = parser.parse_args()
|
|
|
|
# handle both "AAPL MSFT" and "AAPL,MSFT"
|
|
symbols = []
|
|
for tok in args.symbols:
|
|
symbols.extend(s.strip().upper() for s in tok.split(",") if s.strip())
|
|
|
|
download(symbols, args.start, args.end, args.outdir)
|
|
print("Done.")
|