128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
import asyncio
|
|
import logging
|
|
import re
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
import pandas as pd
|
|
from ib_insync import IB, Contract
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_UNIT_SECONDS = {
|
|
"sec": 1,
|
|
"min": 60,
|
|
"hour": 3600,
|
|
"day": 86400,
|
|
"week": 604800,
|
|
"month": 2592000,
|
|
}
|
|
|
|
def parse_bar_size_seconds(bar_size: str) -> int:
|
|
"""Parse IB bar size strings like '1 min', '5 secs', '1 hour' to seconds."""
|
|
m = re.match(r"^\s*(\d+)\s*([a-zA-Z]+)\s*$", bar_size)
|
|
if not m:
|
|
raise ValueError(f"Unsupported bar size: {bar_size!r}")
|
|
n, unit = int(m.group(1)), m.group(2).lower()
|
|
for prefix, seconds in _UNIT_SECONDS.items():
|
|
if unit.startswith(prefix):
|
|
return n * seconds
|
|
raise ValueError(f"Unsupported bar size unit: {bar_size!r}")
|
|
|
|
class BarManager:
|
|
"""Maintains one live (keepUpToDate) historical bar subscription per
|
|
(contract, bar size, data type), shared by all strategies.
|
|
|
|
This replaces per-cycle full history downloads: after the initial fetch,
|
|
IB pushes incremental updates, cutting traffic dramatically.
|
|
"""
|
|
|
|
def __init__(self, ib: IB):
|
|
self.ib = ib
|
|
# key -> BarDataList (ready) or asyncio.Task (subscription in flight).
|
|
# Storing the in-flight task makes concurrent subscribe() calls for the
|
|
# same key share one request instead of firing duplicate subscriptions.
|
|
self._bars: dict[tuple, object] = {}
|
|
|
|
@staticmethod
|
|
def _key(contract: Contract, bar_size: str, what_to_show: str) -> tuple:
|
|
return (contract.conId, bar_size, what_to_show)
|
|
|
|
async def subscribe(self, contract: Contract, bar_size: str, duration: str,
|
|
what_to_show: str, use_rth: bool = True):
|
|
|
|
key = self._key(contract, bar_size, what_to_show)
|
|
existing = self._bars.get(key)
|
|
if existing is not None:
|
|
if isinstance(existing, asyncio.Task):
|
|
return await existing # share the in-flight request
|
|
return existing
|
|
|
|
task = asyncio.create_task(self.ib.reqHistoricalDataAsync(
|
|
contract=contract,
|
|
endDateTime="",
|
|
durationStr=duration,
|
|
barSizeSetting=bar_size,
|
|
whatToShow=what_to_show,
|
|
useRTH=use_rth,
|
|
formatDate=2, # epoch seconds -> unambiguous UTC timestamps
|
|
keepUpToDate=True,
|
|
))
|
|
self._bars[key] = task
|
|
try:
|
|
bars = await task
|
|
except Exception:
|
|
self._bars.pop(key, None) # allow retry on failure
|
|
raise
|
|
self._bars[key] = bars
|
|
logger.info(
|
|
"Bar subscription started: conId=%s %s %s (%d bars)",
|
|
contract.conId, bar_size, what_to_show, len(bars),
|
|
)
|
|
return bars
|
|
|
|
def get(self, contract: Contract, bar_size: str, what_to_show: str):
|
|
|
|
v = self._bars.get(self._key(contract, bar_size, what_to_show))
|
|
return None if isinstance(v, asyncio.Task) else v
|
|
|
|
def reset(self):
|
|
"""Cancel all live subscriptions (e.g. after a disconnect)."""
|
|
|
|
for key, v in self._bars.items():
|
|
try:
|
|
if isinstance(v, asyncio.Task):
|
|
v.cancel()
|
|
else:
|
|
self.ib.cancelHistoricalData(v)
|
|
except Exception:
|
|
pass
|
|
if self._bars:
|
|
logger.info("Bar subscriptions cleared (%d)", len(self._bars))
|
|
self._bars.clear()
|
|
|
|
def to_completed_df(bars, bar_seconds: int) -> pd.DataFrame:
|
|
"""Convert bars to a DataFrame containing only completed bars.
|
|
|
|
The last bar of a live subscription is still forming; signals computed on
|
|
it would repaint, so it is dropped until it closes.
|
|
"""
|
|
if not bars:
|
|
return pd.DataFrame()
|
|
df = pd.DataFrame(bars)
|
|
df["date"] = pd.to_datetime(df["date"], utc=True)
|
|
now = datetime.now(timezone.utc)
|
|
if len(df) and df["date"].iloc[-1] + timedelta(seconds=bar_seconds) > now:
|
|
df = df.iloc[:-1]
|
|
return df
|
|
|
|
def last_bar_age_seconds(df: pd.DataFrame) -> float:
|
|
"""Age of the most recent completed bar. Large age => market closed/stale."""
|
|
if df.empty:
|
|
return float("inf")
|
|
return (datetime.now(timezone.utc) - df["date"].iloc[-1]).total_seconds()
|
|
|
|
def is_market_active(df: pd.DataFrame, bar_seconds: int, factor: int = 5) -> bool:
|
|
"""Data-driven market-hours check: if the latest completed bar is too old,
|
|
the market is closed (overnight/weekend/holiday) and we must not trade."""
|
|
return last_bar_age_seconds(df) <= bar_seconds * factor
|