lumibot 4.1.3__py3-none-any.whl → 4.2.1__py3-none-any.whl

lumibot/data_sources/yahoo_data_polars.py

@@ -7,6 +7,11 @@ This implementation:
 4. Efficient caching with parquet files
 5. Vectorized operations only
 """
+# NOTE: This module is intentionally disabled. The DataBento Polars migration only
+# supports Polars for DataBento; other data sources must use the pandas implementations.
+raise RuntimeError('Yahoo/Polygon Polars backends are not production-ready; use the pandas data sources instead.')
+
+
 
 from datetime import datetime, timedelta
 from decimal import Decimal

lumibot/entities/__init__.py

@@ -5,6 +5,7 @@ from .bar import Bar
 from .bars import Bars as _BarsBase
 from .chains import Chains
 from .data import Data as _DataBase
+from .data_polars import DataPolars
 from .dataline import Dataline
 from .order import Order
 from .position import Position
@@ -14,3 +15,17 @@ from .trading_fee import TradingFee
 # Use base implementations directly
 Bars = _BarsBase
 Data = _DataBase
+__all__ = [
+    "Asset",
+    "AssetsMapping",
+    "Bar",
+    "Bars",
+    "Chains",
+    "Data",
+    "DataPolars",
+    "Dataline",
+    "Order",
+    "Position",
+    "Quote",
+    "TradingFee",
+]

lumibot/entities/asset.py

@@ -715,35 +715,12 @@ class Asset:
         # logger = logging.getLogger(__name__)
         # logger.info(f"[CONTRACT RESOLUTION] symbol={self.symbol}, reference_date={reference_date}, month={reference_date.month}, day={reference_date.day}")
 
-        current_month = reference_date.month
-        current_year = reference_date.year
-        current_day = reference_date.day
+        from lumibot.tools import futures_roll
 
-        # Use quarterly contracts (Mar, Jun, Sep, Dec) with a mid-month roll rule
-        if current_month == 12 and current_day >= 15:
-            target_month = 3
-            target_year = current_year + 1
-        elif current_month >= 10:
-            target_month = 12
-            target_year = current_year
-        elif current_month == 9 and current_day >= 15:
-            target_month = 12
-            target_year = current_year
-        elif current_month >= 7:
-            target_month = 9
-            target_year = current_year
-        elif current_month == 6 and current_day >= 15:
-            target_month = 9
-            target_year = current_year
-        elif current_month >= 4:
-            target_month = 6
-            target_year = current_year
-        elif current_month == 3 and current_day >= 15:
-            target_month = 6
-            target_year = current_year
-        else:
-            target_month = 3
-            target_year = current_year
+        target_year, target_month = futures_roll.determine_contract_year_month(
+            self.symbol,
+            reference_date,
+        )
 
         month_code = FUTURES_MONTH_CODES.get(target_month, "U")
         base_contract = f"{self.symbol}{month_code}"

lumibot/entities/bars.py

@@ -179,8 +179,9 @@ class Bars:
         self.quote = quote
         self._raw = raw
         self._return_polars = return_polars
-        # Cache for on-demand conversions to avoid repeated expensive copies
+        # Caches for on-demand conversions to avoid repeated expensive copies
         self._polars_cache = None
+        self._pandas_cache = None
         self._tzinfo = self._normalize_tzinfo(tzinfo)
         
         # Check if empty
@@ -204,9 +205,20 @@ class Bars:
                     pl.col("close").pct_change().alias("return")
                 ])
             
+            if "datetime" in df.columns and self._tzinfo is not None:
+                target_tz = getattr(self._tzinfo, "zone", None) or getattr(self._tzinfo, "key", None)
+                if target_tz:
+                    current_dtype = df.schema.get("datetime")
+                    current_tz = getattr(current_dtype, "time_zone", None)
+                    if current_tz != target_tz:
+                        df = df.with_columns(
+                            pl.col("datetime").dt.convert_time_zone(target_tz)
+                        )
+
             if return_polars:
                 # Keep as polars
                 self._df = df
+                self._pandas_cache = None
             else:
                 # Convert to pandas and track the conversion
                 tracker = PolarsConversionTracker()
@@ -219,6 +231,7 @@ class Bars:
                         self._df = self._df.set_index(col_name)
                         break
                 self._apply_timezone()
+                self._pandas_cache = None
         else:
             # Already pandas, keep it as is
             self._df = df
@@ -231,19 +244,38 @@ class Bars:
                 self._df["return"] = df["close"].pct_change()
 
             self._apply_timezone()
+            if self._return_polars:
+                self._pandas_cache = self._df
+                self._df = self._convert_pandas_to_polars(self._df)
+                self._polars_cache = None
 
     @property
     def df(self):
-        """Return the DataFrame"""
-        return self._df
+        """Return the active DataFrame representation based on return_polars flag."""
+        return self.polars_df if self._return_polars else self.pandas_df
 
     @df.setter
     def df(self, value):
-        """Allow setting the DataFrame"""
+        """Allow setting the DataFrame while keeping caches in sync."""
         self._df = value
-        # Invalidate cached converted forms when df changes
         self._polars_cache = None
-        self._apply_timezone()
+        self._pandas_cache = None
+
+        if isinstance(value, pd.DataFrame):
+            self._apply_timezone()
+            if self._return_polars:
+                self._pandas_cache = value
+                self._df = self._convert_pandas_to_polars(value)
+        elif isinstance(value, pl.DataFrame) and not self._return_polars:
+            tracker = PolarsConversionTracker()
+            tracker.track_conversion(self.asset.symbol if hasattr(self.asset, 'symbol') else str(self.asset))
+            pandas_df = value.to_pandas()
+            for col_name in ['datetime', 'timestamp', 'date', 'time']:
+                if col_name in pandas_df.columns:
+                    pandas_df = pandas_df.set_index(col_name)
+                    break
+            self._df = pandas_df
+            self._apply_timezone()
 
     @property
     def polars_df(self):
@@ -254,12 +286,29 @@ class Bars:
             # Convert pandas to polars once and cache
             if self._polars_cache is not None:
                 return self._polars_cache
-            if hasattr(self._df, 'index') and getattr(self._df.index, 'name', None):
-                self._polars_cache = pl.from_pandas(self._df.reset_index())
-            else:
-                self._polars_cache = pl.from_pandas(self._df)
+            self._polars_cache = self._convert_pandas_to_polars(self._df)
             return self._polars_cache
 
+    @property
+    def pandas_df(self):
+        """Return as Pandas DataFrame, converting on demand if required."""
+        if isinstance(self._df, pd.DataFrame):
+            return self._df
+        if self._pandas_cache is not None:
+            return self._pandas_cache
+
+        tracker = PolarsConversionTracker()
+        tracker.track_conversion(self.asset.symbol if hasattr(self.asset, 'symbol') else str(self.asset))
+
+        pandas_df = self._df.to_pandas()
+        for col_name in ['datetime', 'timestamp', 'date', 'time']:
+            if col_name in pandas_df.columns:
+                pandas_df = pandas_df.set_index(col_name)
+                break
+
+        self._pandas_cache = self._apply_timezone(pandas_df)
+        return self._pandas_cache
+
     def __repr__(self):
         return repr(self.df)
 
@@ -268,7 +317,12 @@ class Bars:
 
     def __len__(self):
         """Return the number of bars (rows) in the DataFrame"""
-        return len(self.df)
+        if isinstance(self._df, pl.DataFrame):
+            return self._df.height
+        if isinstance(self._df, pd.DataFrame):
+            return len(self._df)
+        df = self.df
+        return df.height if isinstance(df, pl.DataFrame) else len(df)
 
     @property
     def empty(self):
@@ -285,24 +339,39 @@ class Bars:
             return pytz.timezone(tzinfo)
         return tzinfo
 
-    def _apply_timezone(self):
-        if not isinstance(self._df, pd.DataFrame):
-            return
-        if not isinstance(self._df.index, pd.DatetimeIndex):
-            return
+    def _apply_timezone(self, df=None):
+        target_df = self._df if df is None else df
+        if not isinstance(target_df, pd.DataFrame):
+            return target_df
+        if not isinstance(target_df.index, pd.DatetimeIndex):
+            return target_df
 
         tz = self._tzinfo or LUMIBOT_DEFAULT_PYTZ
         if isinstance(tz, str):
             tz = pytz.timezone(tz)
 
         try:
-            if self._df.index.tz is None:
-                self._df.index = self._df.index.tz_localize(tz)
+            if target_df.index.tz is None:
+                target_df.index = target_df.index.tz_localize(tz)
             else:
-                self._df.index = self._df.index.tz_convert(tz)
+                target_df.index = target_df.index.tz_convert(tz)
             self._tzinfo = tz
         except Exception:
-            pass
+            return target_df
+
+        if df is None:
+            self._df = target_df
+        return target_df
+
+    def _convert_pandas_to_polars(self, pandas_df: pd.DataFrame) -> pl.DataFrame:
+        """Convert a pandas DataFrame (possibly with datetime index) to Polars."""
+        df_to_convert = pandas_df.copy()
+        if isinstance(df_to_convert.index, pd.DatetimeIndex):
+            df_to_convert = df_to_convert.reset_index()
+            first_col = df_to_convert.columns[0]
+            if first_col != "datetime":
+                df_to_convert = df_to_convert.rename(columns={first_col: "datetime"})
+        return pl.from_pandas(df_to_convert)
 
     @classmethod
     def parse_bar_list(cls, bar_list, source, asset):

lumibot/entities/data.py

@@ -411,6 +411,14 @@ class Data:
 
             length = kwargs.get("length", 1)
             timeshift = kwargs.get("timeshift", 0)
+
+            if isinstance(timeshift, datetime.timedelta):
+                if self.timestep == "day":
+                    timeshift = int(timeshift.total_seconds() / (24 * 3600))
+                else:
+                    timeshift = int(timeshift.total_seconds() / 60)
+                kwargs["timeshift"] = timeshift
+
             data_index = i + 1 - length - timeshift
             is_data = data_index >= 0
             if not is_data:
@@ -437,8 +445,8 @@ class Data:
             The number of periods to get the last price.
         timestep : str
             The frequency of the data to get the last price.
-        timeshift : int
-            The number of periods to shift the data.
+        timeshift : int | datetime.timedelta
+            The number of periods to shift the data, or a timedelta that will be converted to periods.
 
         Returns
         -------
@@ -462,8 +470,8 @@ class Data:
             The number of periods to get the last price.
         timestep : str
             The frequency of the data to get the last price.
-        timeshift : int
-            The number of periods to shift the data.
+        timeshift : int | datetime.timedelta
+            The number of periods to shift the data, or a timedelta that will be converted to periods.
 
         Returns
         -------
@@ -547,12 +555,27 @@ class Data:
 
         """
 
-        # Get bars.
+        if isinstance(timeshift, datetime.timedelta):
+            if self.timestep == "day":
+                timeshift = int(timeshift.total_seconds() / (24 * 3600))
+            else:
+                timeshift = int(timeshift.total_seconds() / 60)
+
         end_row = self.get_iter_count(dt) - timeshift
-        start_row = end_row - length
 
+        data_len = len(next(iter(self.datalines.values())).dataline) if self.datalines else 0
+        if end_row > data_len:
+            end_row = data_len
+        if end_row < 0:
+            end_row = 0
+
+        start_row = end_row - length
         if start_row < 0:
             start_row = 0
+        if start_row > end_row:
+            start_row = end_row
+        if start_row == end_row and end_row > 0:
+            start_row = max(0, end_row - 1)
 
         # Cast both start_row and end_row to int
         start_row = int(start_row)