lumibot 4.1.2__py3-none-any.whl → 4.2.0__py3-none-any.whl

lumibot/tools/futures_roll.py

@@ -0,0 +1,251 @@
+"""Centralised futures roll logic shared by assets, data sources, and brokers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Dict, List, Optional, Tuple
+
+import pytz
+
+from lumibot.constants import LUMIBOT_DEFAULT_PYTZ
+
+_FUTURES_MONTH_CODES: Dict[int, str] = {
+    1: "F",
+    2: "G",
+    3: "H",
+    4: "J",
+    5: "K",
+    6: "M",
+    7: "N",
+    8: "Q",
+    9: "U",
+    10: "V",
+    11: "X",
+    12: "Z",
+}
+
+
+@dataclass(frozen=True)
+class RollRule:
+    offset_business_days: int
+    anchor: str
+
+
+ROLL_RULES: Dict[str, RollRule] = {
+    symbol: RollRule(offset_business_days=8, anchor="third_friday")
+    for symbol in {"ES", "MES", "NQ", "MNQ", "YM", "MYM"}
+}
+
+YearMonth = Tuple[int, int]
+
+
+def _to_timezone(dt: datetime, tz=pytz.timezone("America/New_York")) -> datetime:
+    if dt.tzinfo is None:
+        return tz.localize(dt)
+    return dt.astimezone(tz)
+
+
+def _normalize_reference_date(reference_date: Optional[datetime]) -> datetime:
+    if reference_date is None:
+        reference_date = datetime.utcnow()
+    return _to_timezone(reference_date, LUMIBOT_DEFAULT_PYTZ)
+
+
+def _third_friday(year: int, month: int) -> datetime:
+    first = datetime(year, month, 1)
+    first = _to_timezone(first)
+    weekday = first.weekday()
+    days_until_friday = (4 - weekday) % 7
+    first_friday = first + timedelta(days=days_until_friday)
+    third_friday = first_friday + timedelta(weeks=2)
+    return third_friday.replace(hour=0, minute=0, second=0, microsecond=0)
+
+
+def _subtract_business_days(dt: datetime, days: int) -> datetime:
+    result = dt
+    remaining = days
+    while remaining > 0:
+        result -= timedelta(days=1)
+        if result.weekday() < 5:
+            remaining -= 1
+    return result
+
+
+def _calculate_roll_trigger(year: int, month: int, rule: RollRule) -> datetime:
+    if rule.anchor == "third_friday":
+        anchor = _third_friday(year, month)
+    else:
+        anchor = _to_timezone(datetime(year, month, 15))
+    if rule.offset_business_days <= 0:
+        return anchor
+    return _subtract_business_days(anchor, rule.offset_business_days)
+
+
+def _advance_quarter(current_month: int, current_year: int) -> YearMonth:
+    quarter_months = [3, 6, 9, 12]
+    idx = quarter_months.index(current_month)
+    next_idx = (idx + 1) % len(quarter_months)
+    next_month = quarter_months[next_idx]
+    next_year = current_year + (1 if next_idx == 0 else 0)
+    return next_year, next_month
+
+
+def _legacy_mid_month(reference_date: datetime) -> YearMonth:
+    quarter_months = [3, 6, 9, 12]
+    year = reference_date.year
+    month = reference_date.month
+    day = reference_date.day
+
+    if month == 12 and day >= 15:
+        return year + 1, 3
+    if month >= 10:
+        return year, 12
+    if month == 9 and day >= 15:
+        return year, 12
+    if month >= 7:
+        return year, 9
+    if month == 6 and day >= 15:
+        return year, 9
+    if month >= 4:
+        return year, 6
+    if month == 3 and day >= 15:
+        return year, 6
+    return year, 3
+
+
+def determine_contract_year_month(symbol: str, reference_date: Optional[datetime] = None) -> YearMonth:
+    ref = _normalize_reference_date(reference_date)
+    symbol_upper = symbol.upper()
+    rule = ROLL_RULES.get(symbol_upper)
+
+    quarter_months = [3, 6, 9, 12]
+    year = ref.year
+    month = ref.month
+
+    if rule is None:
+        return _legacy_mid_month(ref)
+
+    if month in quarter_months:
+        target_year, target_month = year, month
+        roll_point = _calculate_roll_trigger(target_year, target_month, rule)
+        if ref >= roll_point:
+            target_year, target_month = _advance_quarter(target_month, target_year)
+    else:
+        candidates = [m for m in quarter_months if m > month]
+        if candidates:
+            target_month = candidates[0]
+            target_year = year
+        else:
+            target_month = quarter_months[0]
+            target_year = year + 1
+
+    return target_year, target_month
+
+
+def build_contract_symbol(root: str, year: int, month: int, year_digits: int = 2) -> str:
+    month_code = _FUTURES_MONTH_CODES.get(month)
+    if month_code is None:
+        raise ValueError(f"Unsupported futures month: {month}")
+    if year_digits == 1:
+        return f"{root}{month_code}{year % 10}"
+    if year_digits == 4:
+        return f"{root}{month_code}{year}"
+    return f"{root}{month_code}{year % 100:02d}"
+
+
+def resolve_symbol_for_datetime(asset, dt: datetime, year_digits: int = 2) -> str:
+    year, month = determine_contract_year_month(asset.symbol, dt)
+    return build_contract_symbol(asset.symbol, year, month, year_digits=year_digits)
+
+
+def resolve_symbols_for_range(asset, start: datetime, end: datetime, year_digits: int = 2) -> List[str]:
+    if start is None or end is None:
+        return []
+
+    start = _normalize_reference_date(start)
+    end = _normalize_reference_date(end)
+    if start > end:
+        start, end = end, start
+
+    symbols: List[str] = []
+    seen: set[str] = set()
+    cursor = start
+    step = timedelta(days=30)
+
+    while cursor <= end + timedelta(days=45):
+        symbol = resolve_symbol_for_datetime(asset, cursor, year_digits=year_digits)
+        if symbol not in seen:
+            seen.add(symbol)
+            symbols.append(symbol)
+        cursor += step
+
+    final_symbol = resolve_symbol_for_datetime(asset, end, year_digits=year_digits)
+    if final_symbol not in seen:
+        symbols.append(final_symbol)
+
+    if final_symbol in symbols:
+        final_index = symbols.index(final_symbol)
+        symbols = symbols[: final_index + 1]
+
+    return symbols
+
+def build_roll_schedule(asset, start: datetime, end: datetime, year_digits: int = 2):
+    if start is None or end is None:
+        return []
+
+    start = _normalize_reference_date(start)
+    end = _normalize_reference_date(end)
+    if start > end:
+        start, end = end, start
+
+    symbol_upper = asset.symbol.upper()
+    rule = ROLL_RULES.get(symbol_upper)
+
+    schedule = []
+    cursor = start
+    previous_start = start
+
+    while cursor <= end + timedelta(days=90):
+        year, month = determine_contract_year_month(symbol_upper, cursor)
+        symbol = build_contract_symbol(symbol_upper, year, month, year_digits=year_digits)
+
+        if rule:
+            roll_dt = _calculate_roll_trigger(year, month, rule)
+        else:
+            roll_dt = _to_timezone(datetime(year, month, 15))
+
+        schedule.append((symbol, previous_start, roll_dt))
+
+        cursor = roll_dt + timedelta(minutes=1)
+        previous_start = cursor
+        if roll_dt >= end:
+            break
+
+    clipped = []
+    for symbol, s, e in schedule:
+        start_clip = max(s, start)
+        end_clip = min(e, end)
+        if end_clip <= start_clip:
+            continue
+        clipped.append((symbol, start_clip, end_clip))
+
+    if not clipped:
+        return [(
+            symbol,
+            s.astimezone(pytz.UTC),
+            e.astimezone(pytz.UTC),
+        ) for symbol, s, e in schedule]
+
+    last_symbol, s, e = clipped[-1]
+    if e < end:
+        clipped[-1] = (last_symbol, s, end)
+
+    return [
+        (
+            symbol,
+            start_clip.astimezone(pytz.UTC),
+            end_clip.astimezone(pytz.UTC),
+        )
+        for symbol, start_clip, end_clip in clipped
+    ]

lumibot/tools/indicators.py

@@ -18,8 +18,142 @@ from plotly.subplots import make_subplots
 from .yahoo_helper import YahooHelper as yh
 
 from lumibot.tools.lumibot_logger import get_logger
+
 logger = get_logger(__name__)
 
+TERMINAL_TRADE_STATUSES_FOR_MARKERS = {
+    "fill",
+    "filled",
+    "partial_fill",
+    "cash_settled",
+    "assigned",
+    "assignment",
+    "exercise",
+    "exercised",
+    "expired",
+    "expire",
+}
+
+
+def _build_trade_marker_tooltip(row: pd.Series):
+    """Return tooltip text for a trade marker; None when the row lacks required data."""
+    status_value = row.get("status")
+    if pd.isna(status_value) or str(status_value).strip() == "":
+        return None
+
+    status_text = str(status_value)
+    if status_text.lower() not in TERMINAL_TRADE_STATUSES_FOR_MARKERS:
+        return None
+
+    for key in ("filled_quantity", "price"):
+        value = row.get(key)
+        if pd.isna(value):
+            return None
+
+    try:
+        filled_quantity_dec = Decimal(str(row["filled_quantity"]))
+        price_dec = Decimal(str(row["price"]))
+    except (InvalidOperation, TypeError, ValueError):
+        return None
+
+    multiplier_value = row.get("asset.multiplier")
+    if pd.isna(multiplier_value) or multiplier_value == "":
+        return None
+    try:
+        multiplier_dec = Decimal(str(multiplier_value))
+    except (InvalidOperation, TypeError, ValueError):
+        return None
+
+    try:
+        amount_transacted_dec = price_dec * filled_quantity_dec * multiplier_dec
+    except (InvalidOperation, TypeError, ValueError):
+        return None
+
+    trade_cost_value = row.get("trade_cost")
+    trade_cost_dec = None
+    if not (pd.isna(trade_cost_value) or trade_cost_value == ""):
+        try:
+            trade_cost_dec = Decimal(str(trade_cost_value))
+        except (InvalidOperation, TypeError, ValueError):
+            trade_cost_dec = None
+
+    if trade_cost_dec is None:
+        trade_cost_dec = amount_transacted_dec
+
+    if row.get("asset.asset_type") == "option":
+        try:
+            return (
+                status_text
+                + "<br>"
+                + str(filled_quantity_dec.quantize(Decimal("0.01")).__format__(",f"))
+                + " "
+                + str(row.get("symbol"))
+                + " "
+                + str(row.get("asset.right"))
+                + " Option"
+                + "<br>"
+                + "Strike: "
+                + str(row.get("asset.strike"))
+                + "<br>"
+                + "Expiration: "
+                + str(row.get("asset.expiration"))
+                + "<br>"
+                + "Price: "
+                + str(price_dec.quantize(Decimal("0.0001")).__format__(",f"))
+                + "<br>"
+                + "Order Type: "
+                + str(row.get("type"))
+                + "<br>"
+                + "Amount Transacted: "
+                + str(
+                    (
+                        price_dec
+                        * filled_quantity_dec
+                        * (multiplier_dec if multiplier_dec != Decimal("0") else Decimal("1"))
+                    )
+                    .quantize(Decimal("0.01"))
+                    .__format__(",f")
+                )
+                + "<br>"
+                + "Trade Cost: "
+                + str(trade_cost_dec.quantize(Decimal("0.01")).__format__(",f"))
+                + "<br>"
+            )
+        except (InvalidOperation, TypeError, ValueError):
+            return None
+
+    if multiplier_dec == Decimal("0"):
+        return None
+
+    try:
+        amount_transacted = amount_transacted_dec.quantize(Decimal("0.01")).__format__(",f")
+        price_text = str(price_dec.quantize(Decimal("0.0001")).__format__(",f"))
+        filled_qty_text = str(filled_quantity_dec.quantize(Decimal("0.01")).__format__(",f"))
+        trade_cost_text = str(trade_cost_dec.quantize(Decimal("0.01")).__format__(",f"))
+    except (InvalidOperation, TypeError, ValueError):
+        return None
+
+    return (
+        status_text
+        + "<br>"
+        + filled_qty_text
+        + " "
+        + str(row.get("symbol"))
+        + "<br>"
+        + "Price: "
+        + price_text
+        + "<br>"
+        + "Order Type: "
+        + str(row.get("type"))
+        + "<br>"
+        + "Amount Transacted: "
+        + amount_transacted
+        + "<br>"
+        + "Trade Cost: "
+        + trade_cost_text
+        + "<br>"
+    )
+
 
 def total_return(_df):
     """Calculate the cumulative return in a dataframe
@@ -669,110 +803,7 @@ def plot_returns(
     buys = buys.loc[df_final["side"].isin(["buy", "buy_to_open", "buy_to_cover", "buy_to_close"])]
 
     def generate_buysell_plotly_text(row):
-        if row["status"] not in ("fill", "partial_fill"):
-            return None
-
-        for key in ("filled_quantity", "price"):
-            value = row.get(key)
-            if pd.isna(value):
-                return None
-
-        try:
-            filled_quantity_dec = Decimal(str(row["filled_quantity"]))
-            price_dec = Decimal(str(row["price"]))
-        except (InvalidOperation, TypeError, ValueError):
-            return None
-
-        multiplier_value = row.get("asset.multiplier")
-        if pd.isna(multiplier_value) or multiplier_value == "":
-            return None
-        try:
-            multiplier_dec = Decimal(str(multiplier_value))
-        except (InvalidOperation, TypeError, ValueError):
-            return None
-
-        trade_cost_value = row.get("trade_cost")
-        if pd.isna(trade_cost_value) or trade_cost_value == "":
-            return None
-        try:
-            trade_cost_dec = Decimal(str(trade_cost_value))
-        except (InvalidOperation, TypeError, ValueError):
-            return None
-
-        if row["asset.asset_type"] == "option":
-            try:
-                return (
-                    row["status"]
-                    + "<br>"
-                    + str(filled_quantity_dec.quantize(Decimal("0.01")).__format__(",f"))
-                    + " "
-                    + row["symbol"]
-                    + " "
-                    + row["asset.right"]
-                    + " Option"
-                    + "<br>"
-                    + "Strike: "
-                    + str(row["asset.strike"])
-                    + "<br>"
-                    + "Expiration: "
-                    + str(row["asset.expiration"])
-                    + "<br>"
-                    + "Price: "
-                    + str(price_dec.quantize(Decimal("0.0001")).__format__(",f"))
-                    + "<br>"
-                    + "Order Type: "
-                    + row["type"]
-                    + "<br>"
-                    + "Amount Transacted: "
-                    + str(
-                        (
-                            price_dec
-                            * filled_quantity_dec
-                            * (multiplier_dec if multiplier_dec != Decimal("0") else Decimal("1"))
-                        )
-                        .quantize(Decimal("0.01"))
-                        .__format__(",f")
-                    )
-                    + "<br>"
-                    + "Trade Cost: "
-                    + str(trade_cost_dec.quantize(Decimal("0.01")).__format__(",f"))
-                    + "<br>"
-                )
-            except (InvalidOperation, TypeError, ValueError):
-                return None
-
-        if multiplier_dec == Decimal("0"):
-            return None
-        try:
-            amount_transacted = (
-                price_dec * filled_quantity_dec * multiplier_dec
-            ).quantize(Decimal("0.01")).__format__(",f")
-            price_text = str(price_dec.quantize(Decimal("0.0001")).__format__(",f"))
-            filled_qty_text = str(filled_quantity_dec.quantize(Decimal("0.01")).__format__(",f"))
-            trade_cost_text = str(trade_cost_dec.quantize(Decimal("0.01")).__format__(",f"))
-        except (InvalidOperation, TypeError, ValueError):
-            return None
-
-        return (
-            row["status"]
-            + "<br>"
-            + filled_qty_text
-            + " "
-            + row["symbol"]
-            + "<br>"
-            + "Price: "
-            + price_text
-            + "<br>"
-            + "Order Type: "
-            + row["type"]
-            + "<br>"
-            + "Amount Transacted: "
-            + amount_transacted
-            + "<br>"
-            + "Trade Cost: "
-            + trade_cost_text
-            + "<br>"
-        )
+        return _build_trade_marker_tooltip(row)
 
     buy_ticks_df = buys.apply(generate_buysell_plotly_text, axis=1)
 

lumibot/tools/polars_utils.py

@@ -0,0 +1,142 @@
+"""Utility helpers for operating on Polars DataFrames within Lumibot."""
+
+from __future__ import annotations
+
+from typing import Iterable, Optional, Sequence, Set
+
+import polars as pl
+
+
+class PolarsResampleError(Exception):
+    """Raised when a Polars resample operation cannot be completed."""
+
+
+def _ensure_datetime_column(df: pl.DataFrame) -> str:
+    """Return the datetime-like column name used for grouping."""
+    if "datetime" in df.columns:
+        return "datetime"
+
+    for candidate in ("timestamp", "date", "time"):
+        if candidate in df.columns:
+            return candidate
+
+    raise PolarsResampleError("Polars DataFrame lacks a datetime-like column required for resampling.")
+
+
+def _aggregate_expressions(existing_cols: Sequence[str]) -> list[pl.Expr]:
+    """Build aggregation expressions for OHLC-style resampling."""
+    exprs: list[pl.Expr] = []
+    handled: Set[str] = {"datetime", "timestamp", "date", "time"}
+
+    if "open" in existing_cols:
+        exprs.append(pl.col("open").first().alias("open"))
+        handled.add("open")
+
+    if "high" in existing_cols:
+        exprs.append(pl.col("high").max().alias("high"))
+        handled.add("high")
+
+    if "low" in existing_cols:
+        exprs.append(pl.col("low").min().alias("low"))
+        handled.add("low")
+
+    if "close" in existing_cols:
+        exprs.append(pl.col("close").last().alias("close"))
+        handled.add("close")
+
+    if "volume" in existing_cols:
+        exprs.append(pl.col("volume").sum().alias("volume"))
+        handled.add("volume")
+
+    if "dividend" in existing_cols:
+        exprs.append(pl.col("dividend").sum().alias("dividend"))
+        handled.add("dividend")
+
+    # Preserve any remaining columns by taking the last observation
+    for column in existing_cols:
+        if column not in handled:
+            exprs.append(pl.col(column).last().alias(column))
+
+    return exprs
+
+
+def resample_polars_ohlc(
+    df: pl.DataFrame,
+    multiplier: int,
+    base_unit: str,
+    length: Optional[int] = None,
+    label_offset: Optional[str] = None,
+) -> pl.DataFrame:
+    """Resample a Polars DataFrame containing OHLC-like data.
+
+    Parameters
+    ----------
+    df:
+        Input DataFrame containing at least ``datetime`` plus OHLCV columns.
+    multiplier:
+        Number of base units to roll up. e.g. multiplier=5, base_unit="minute" -> 5-minute bars.
+    base_unit:
+        Currently supports "minute" or "day".
+    length:
+        Optional maximum number of rows to retain (tail). If ``None`` retains the full frame.
+    label_offset:
+        Optional duration string understood by Polars to offset labels. Useful for aligning session boundaries.
+
+    Returns
+    -------
+    pl.DataFrame
+        Resampled dataset sorted by datetime.
+    """
+
+    if df.is_empty():
+        return df
+
+    if multiplier <= 0:
+        raise PolarsResampleError("Multiplier must be positive for resampling.")
+
+    unit_map = {"minute": "m", "day": "d"}
+    try:
+        every_suffix = unit_map[base_unit]
+    except KeyError as exc:
+        raise PolarsResampleError(f"Unsupported base unit '{base_unit}' for polars resampling.") from exc
+
+    every = f"{multiplier}{every_suffix}"
+
+    datetime_column = _ensure_datetime_column(df)
+    sorted_df = df.sort(datetime_column)
+
+    agg_exprs = _aggregate_expressions(sorted_df.columns)
+
+    group_kwargs = {
+        "every": every,
+        "period": every,
+        "closed": "left",
+        "label": "left",
+    }
+    if label_offset:
+        group_kwargs["offset"] = label_offset
+
+    lazy_frame = sorted_df.lazy()
+    if hasattr(lazy_frame, "group_by_dynamic"):
+        lazy_grouped = lazy_frame.group_by_dynamic(datetime_column, **group_kwargs)
+    else:  # pragma: no cover - backward compatibility
+        lazy_grouped = lazy_frame.groupby_dynamic(datetime_column, **group_kwargs)
+    resampled = (
+        lazy_grouped
+        .agg(agg_exprs)
+        .sort(datetime_column)
+        .collect()
+    )
+
+    required_cols: Iterable[str] = [c for c in ("open", "high", "low", "close") if c in resampled.columns]
+    if required_cols:
+        condition = None
+        for col in required_cols:
+            expr = pl.col(col).is_not_null()
+            condition = expr if condition is None else condition & expr
+        resampled = resampled.filter(condition)
+
+    if length is not None and length > 0 and resampled.height > length:
+        resampled = resampled.tail(length)
+
+    return resampled