Data fetcher implementation

src/trading_bot/data/fetch_ibkr_daily.py

@@ -1,18 +1,20 @@
-"""Skeleton IBKR daily candle fetcher.
-
-This intentionally uses hard-coded values while the project shape is still
-being designed. Parquet persistence and configurable date ranges come later.
-"""
+"""IBKR daily candle fetcher."""
 
 from __future__ import annotations
 
+import argparse
+import re
 from dataclasses import dataclass
-from datetime import date
+from datetime import date, datetime, timedelta
+from pathlib import Path
+from zoneinfo import ZoneInfo
 
+import pandas as pd
 
-TICKER = "SPY"
-START_DATE = date(2026, 6, 1)
-END_DATE = date(2026, 6, 5)
+DEFAULT_OUTPUT_DIR = Path("data/ibkr/daily")
+DEFAULT_DURATION = "1 W"
+EASTERN_TZ = ZoneInfo("America/New_York")
+DURATION_PATTERN = re.compile(r"^\d+\s+[SDWMY]$")
 
 IBKR_HOST = "127.0.0.1"
 IBKR_PORT = 4002
@@ -32,15 +34,67 @@ class DailyCandle:
     volume: int
 
 
-def fetch_hard_coded_daily_candles() -> list[DailyCandle]:
-    """Fetch one hard-coded week of SPY daily candles from IBKR."""
+def default_end_date() -> date:
+    """Return yesterday's date in the US/Eastern market timezone."""
+
+    return datetime.now(EASTERN_TZ).date() - timedelta(days=1)
+
+
+def current_market_date() -> date:
+    """Return today's date in the US/Eastern market timezone."""
+
+    return datetime.now(EASTERN_TZ).date()
+
+
+def parse_end_date(value: str) -> date:
+    """Parse an end date in IBKR-friendly YYYYMMDD form."""
+
+    try:
+        return datetime.strptime(value, "%Y%m%d").date()
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(
+            "end date must use YYYYMMDD format, such as 20250605"
+        ) from exc
+
+
+def parse_duration(value: str) -> str:
+    """Normalize and validate an IBKR duration string like '1 W' or '1 M'."""
+
+    normalized = " ".join(value.upper().split())
+    if not DURATION_PATTERN.fullmatch(normalized):
+        raise argparse.ArgumentTypeError(
+            "duration must look like an IBKR value, such as '1 W' or '1 M'"
+        )
+    return normalized
+
+
+def format_ibkr_end_datetime(end_date: date) -> str:
+    """Convert a date to the US/Eastern end datetime string IBKR expects."""
+
+    return f"{end_date:%Y%m%d} 23:59:59 US/Eastern"
+
+
+def normalize_bar_date(value: date | datetime | str) -> date:
+    """Normalize an IBKR historical bar date value."""
+
+    if isinstance(value, datetime):
+        return value.date()
+    if isinstance(value, date):
+        return value
+    return parse_end_date(value)
+
+
+def fetch_daily_candles(
+    symbol: str, end_date: date, duration: str
+) -> list[DailyCandle]:
+    """Fetch daily candles for a symbol from IBKR."""
 
     try:
         from ib_insync import IB, Stock  # type: ignore[import-not-found]
     except ImportError as exc:
         raise SystemExit(
             "Missing dependency: ib_insync. Install it before running the "
-            "IBKR fetcher skeleton."
+            "IBKR fetcher."
         ) from exc
 
     ib = IB()
@@ -56,12 +110,12 @@ def fetch_hard_coded_daily_candles() -> list[DailyCandle]:
             timeout=IBKR_CONNECT_TIMEOUT_SECONDS,
         )
 
-        contract = Stock(TICKER, "SMART", "USD")
+        contract = Stock(symbol, "SMART", "USD")
         ib.qualifyContracts(contract)
         bars = ib.reqHistoricalData(
             contract,
-            endDateTime="20260605 23:59:59 US/Eastern",
-            durationStr="1 W",
+            endDateTime=format_ibkr_end_datetime(end_date),
+            durationStr=duration,
             barSizeSetting="1 day",
             whatToShow="TRADES",
             useRTH=True,
@@ -69,22 +123,21 @@ def fetch_hard_coded_daily_candles() -> list[DailyCandle]:
         )
 
         if not bars:
-            print("IBKR returned no historical bars for the hard-coded request.")
+            print(f"IBKR returned no historical bars for {symbol}.")
 
         candles: list[DailyCandle] = []
         for bar in bars:
-            trading_day = bar.date if isinstance(bar.date, date) else bar.date.date()
-            if START_DATE <= trading_day <= END_DATE:
-                candles.append(
-                    DailyCandle(
-                        trading_day=trading_day,
-                        open=float(bar.open),
-                        high=float(bar.high),
-                        low=float(bar.low),
-                        close=float(bar.close),
-                        volume=int(bar.volume),
-                    )
+            trading_day = normalize_bar_date(bar.date)
+            candles.append(
+                DailyCandle(
+                    trading_day=trading_day,
+                    open=float(bar.open),
+                    high=float(bar.high),
+                    low=float(bar.low),
+                    close=float(bar.close),
+                    volume=int(bar.volume),
                 )
+            )
 
         return candles
     finally:
@@ -92,14 +145,83 @@ def fetch_hard_coded_daily_candles() -> list[DailyCandle]:
             ib.disconnect()
 
 
-def print_candles(candles: list[DailyCandle]) -> None:
+def candles_to_frame(symbol: str, candles: list[DailyCandle]) -> pd.DataFrame:
+    """Convert candles to a date-indexed dataframe ready for Parquet storage."""
+
+    rows = [
+        {
+            "date": candle.trading_day,
+            "symbol": symbol,
+            "open": candle.open,
+            "high": candle.high,
+            "low": candle.low,
+            "close": candle.close,
+            "volume": candle.volume,
+        }
+        for candle in candles
+    ]
+    frame = pd.DataFrame.from_records(rows)
+    if frame.empty:
+        return pd.DataFrame(
+            columns=["symbol", "open", "high", "low", "close", "volume"],
+            index=pd.Index([], name="date"),
+        )
+
+    frame["date"] = pd.to_datetime(frame["date"]).dt.date
+    return frame.set_index("date")
+
+
+def read_existing_candles(path: Path) -> pd.DataFrame:
+    """Read an existing candle Parquet file as a date-indexed dataframe."""
+
+    if not path.exists():
+        return pd.DataFrame(
+            columns=["symbol", "open", "high", "low", "close", "volume"],
+            index=pd.Index([], name="date"),
+        )
+
+    frame = pd.read_parquet(path)
+    if "date" in frame.columns:
+        frame["date"] = pd.to_datetime(frame["date"]).dt.date
+        frame = frame.set_index("date")
+
+    frame.index = pd.to_datetime(frame.index).date
+    frame.index.name = "date"
+    return frame
+
+
+def oldest_stored_date_or_today(path: Path) -> date:
+    """Return the oldest stored candle date, or today if no data exists yet."""
+
+    existing = read_existing_candles(path)
+    if existing.empty:
+        return current_market_date()
+    return min(existing.index)
+
+
+def write_candles(path: Path, symbol: str, candles: list[DailyCandle]) -> pd.DataFrame:
+    """Append candles to a ticker Parquet file, keeping one row per date."""
+
+    existing = read_existing_candles(path)
+    fetched = candles_to_frame(symbol, candles)
+    combined = pd.concat([existing, fetched])
+    if not combined.empty:
+        combined = combined[~combined.index.duplicated(keep="last")]
+        combined = combined.sort_index()
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    combined.to_parquet(path, index=True)
+    return combined
+
+
+def print_candles(symbol: str, candles: list[DailyCandle]) -> None:
     """Print candles in a compact table."""
 
-    print("date,ticker,open,high,low,close,volume")
+    print("date,symbol,open,high,low,close,volume")
     for candle in candles:
         print(
             f"{candle.trading_day.isoformat()},"
-            f"{TICKER},"
+            f"{symbol},"
             f"{candle.open:.2f},"
             f"{candle.high:.2f},"
             f"{candle.low:.2f},"
@@ -108,12 +230,63 @@ def print_candles(candles: list[DailyCandle]) -> None:
         )
 
 
-def main() -> None:
-    """Run the hard-coded fetcher skeleton."""
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments."""
 
-    print(f"Fetching {TICKER} daily candles from {START_DATE} to {END_DATE}")
-    candles = fetch_hard_coded_daily_candles()
-    print_candles(candles)
+    parser = argparse.ArgumentParser(description="Fetch daily IBKR candles.")
+    parser.add_argument("symbol", help="Ticker symbol to fetch, such as SPY.")
+    end_date_group = parser.add_mutually_exclusive_group()
+    end_date_group.add_argument(
+        "--end-date",
+        type=parse_end_date,
+        help="Request end date in YYYYMMDD format. Defaults to yesterday.",
+    )
+    end_date_group.add_argument(
+        "--end-date-from-parquet",
+        action="store_true",
+        help=(
+            "Use the oldest date from the symbol Parquet file as the request "
+            "end date. Defaults to today if the file is missing or empty."
+        ),
+    )
+    parser.add_argument(
+        "--duration",
+        type=parse_duration,
+        default=DEFAULT_DURATION,
+        help=(
+            "IBKR duration string, such as '1 W' or '1 M'. "
+            f"Defaults to {DEFAULT_DURATION}."
+        ),
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=DEFAULT_OUTPUT_DIR,
+        help=f"Directory for Parquet files. Defaults to {DEFAULT_OUTPUT_DIR}.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Run the daily candle fetcher."""
+
+    args = parse_args()
+    symbol = args.symbol.upper()
+    output_path = args.output_dir / f"{symbol}.parquet"
+    end_date = (
+        oldest_stored_date_or_today(output_path)
+        if args.end_date_from_parquet
+        else args.end_date or default_end_date()
+    )
+
+    print(
+        f"Fetching {symbol} daily candles ending {end_date:%Y-%m-%d} "
+        f"for duration {args.duration}"
+    )
+    candles = fetch_daily_candles(symbol, end_date, args.duration)
+    print_candles(symbol, candles)
+    stored = write_candles(output_path, symbol, candles)
+    print(f"Wrote {len(stored)} total daily rows to {output_path}")
 
 
 if __name__ == "__main__":