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

lumibot/tools/databento_helper.py

@@ -9,12 +9,18 @@ from decimal import Decimal
 import pandas as pd
 from lumibot import LUMIBOT_CACHE_FOLDER
 from lumibot.entities import Asset
-from lumibot.tools import databento_roll
+from lumibot.tools import futures_roll
+from termcolor import colored
 
 # Set up module-specific logger
 from lumibot.tools.lumibot_logger import get_logger
 logger = get_logger(__name__)
 
+
+class DataBentoAuthenticationError(RuntimeError):
+    """Raised when DataBento rejects authentication credentials."""
+    pass
+
 # DataBento imports (will be installed as dependency)
 try:
     import databento as db
@@ -161,11 +167,16 @@ class DataBentoClient:
                         continue
                     else:
                         logger.error(f"DataBento authentication failed after {self.max_retries} retries")
+                        raise DataBentoAuthenticationError(
+                            f"DataBento authentication failed after {self.max_retries} retries: {str(e)}"
+                        ) from e
                         
                 # For non-auth errors, don't retry - fail fast
-                logger.error("DATABENTO_API_ERROR: DataBento API error: %s | Symbols: %s, Start: %s, End: %s", 
-                            str(e), symbols, start, end)
-                raise e
+                logger.error(
+                    "DATABENTO_API_ERROR: DataBento API error: %s | Symbols: %s, Start: %s, End: %s",
+                    str(e), symbols, start, end
+                )
+                raise
         
         # This should never be reached, but just in case
         raise Exception(f"DataBento request failed after {self.max_retries} retries")
@@ -593,13 +604,29 @@ def _filter_front_month_rows_pandas(
     if df.empty or "symbol" not in df.columns or schedule is None:
         return df
 
+    index_tz = getattr(df.index, "tz", None)
+
+    def _align(ts: datetime | pd.Timestamp | None) -> pd.Timestamp | None:
+        if ts is None:
+            return None
+        ts_pd = pd.Timestamp(ts)
+        if index_tz is None:
+            return ts_pd.tz_localize(None) if ts_pd.tz is not None else ts_pd
+        if ts_pd.tz is None:
+            ts_pd = ts_pd.tz_localize(index_tz)
+        else:
+            ts_pd = ts_pd.tz_convert(index_tz)
+        return ts_pd
+
     mask = pd.Series(False, index=df.index)
     for symbol, start_dt, end_dt in schedule:
         cond = df["symbol"] == symbol
-        if start_dt is not None:
-            cond &= df.index >= start_dt
-        if end_dt is not None:
-            cond &= df.index < end_dt
+        start_aligned = _align(start_dt)
+        end_aligned = _align(end_dt)
+        if start_aligned is not None:
+            cond &= df.index >= start_aligned
+        if end_aligned is not None:
+            cond &= df.index < end_aligned
         mask |= cond
 
     filtered = df.loc[mask]
@@ -783,15 +810,31 @@ def get_price_data_from_databento(
     start_naive = start.replace(tzinfo=None) if start.tzinfo is not None else start
     end_naive = end.replace(tzinfo=None) if end.tzinfo is not None else end
 
-    if asset.asset_type == Asset.AssetType.CONT_FUTURE:
+    roll_asset = asset
+    if asset.asset_type == Asset.AssetType.FUTURE and not asset.expiration:
+        roll_asset = Asset(asset.symbol, Asset.AssetType.CONT_FUTURE)
+
+    if roll_asset.asset_type == Asset.AssetType.CONT_FUTURE:
         schedule_start = start
-        symbols = databento_roll.resolve_symbols_for_range(asset, schedule_start, end)
-        front_symbol = databento_roll.resolve_symbol_for_datetime(asset, reference_date or start)
+        symbols = futures_roll.resolve_symbols_for_range(
+            roll_asset,
+            schedule_start,
+            end,
+            year_digits=1,
+        )
+        front_symbol = futures_roll.resolve_symbol_for_datetime(
+            roll_asset,
+            reference_date or start,
+            year_digits=1,
+        )
         if front_symbol not in symbols:
             symbols.insert(0, front_symbol)
     else:
         schedule_start = start
-        front_symbol = _format_futures_symbol_for_databento(asset)
+        front_symbol = _format_futures_symbol_for_databento(
+            asset,
+            reference_date=reference_date or start,
+        )
         symbols = [front_symbol]
 
     # Ensure multiplier is populated using the first contract.
@@ -856,6 +899,13 @@ def get_price_data_from_databento(
                     end=end_naive,
                     **kwargs,
                 )
+            except DataBentoAuthenticationError as exc:
+                auth_msg = colored(
+                    f"❌ DataBento authentication failed while requesting {symbol}: {exc}",
+                    "red"
+                )
+                logger.error(auth_msg)
+                raise
             except Exception as exc:
                 logger.warning(f"Error fetching {symbol} from DataBento: {exc}")
                 continue
@@ -877,38 +927,11 @@ def get_price_data_from_databento(
     combined = pd.concat(frames, axis=0)
     combined.sort_index(inplace=True)
 
-    definition_client: Optional[DataBentoClient] = None
-
-    def get_definition(symbol_code: str) -> Optional[Dict]:
-        nonlocal definition_client
-        cache_key = (symbol_code, dataset)
-        if cache_key in _INSTRUMENT_DEFINITION_CACHE:
-            return _INSTRUMENT_DEFINITION_CACHE[cache_key]
-        if definition_client is None:
-            try:
-                definition_client = DataBentoClient(api_key=api_key)
-            except Exception as exc:
-                logger.warning(f"Unable to create DataBento definition client: {exc}")
-                return None
-        try:
-            definition = definition_client.get_instrument_definition(
-                dataset=dataset,
-                symbol=symbol_code,
-                reference_date=reference_date or start,
-            )
-        except Exception as exc:
-            logger.warning(f"Failed to fetch definition for {symbol_code}: {exc}")
-            return None
-        if definition:
-            _INSTRUMENT_DEFINITION_CACHE[cache_key] = definition
-        return definition
-
-    schedule = databento_roll.build_roll_schedule(
-        asset,
+    schedule = futures_roll.build_roll_schedule(
+        roll_asset,
         schedule_start,
         end,
-        definition_provider=get_definition,
-        roll_days=databento_roll.ROLL_DAYS_BEFORE_EXPIRATION,
+        year_digits=1,
     )
 
     if schedule: