Merge pull request #961 from Lumiwealth/be-fix-get-trading-days

Author: grzesir

lumibot/tools/helpers.py

@@ -160,10 +160,23 @@ def get_trading_days(
 
     # More robust datetime conversion with explicit timezone handling
     def format_datetime(dtm):
+        """
+        Convert dtm to a timezone-aware Python datetime in tzinfo.
+
+        Handles inputs that may be:
+        - tz-naive (localize to tzinfo)
+        - tz-aware (convert to tzinfo)
+        - pandas Timestamp or Python datetime
+        - NaT/None (returned as-is)
+        """
         if pd.isna(dtm):
             return dtm
-        # Convert to Python datetime and ensure proper timezone conversion
-        return pd.Timestamp(dtm).tz_convert(tzinfo).to_pydatetime()
+        ts = pd.Timestamp(dtm)
+        if ts.tz is None:
+            ts = ts.tz_localize(tzinfo)
+        else:
+            ts = ts.tz_convert(tzinfo)
+        return ts.to_pydatetime()
 
     def ensure_tz_aware(dtm, tzinfo):
         dtm = pd.to_datetime(dtm)
@@ -175,15 +188,15 @@ def ensure_tz_aware(dtm, tzinfo):
     else:
         end_date = ensure_tz_aware(get_lumibot_datetime(), tzinfo)
 
-    # Normalize to date-only boundaries (pandas_market_calendars schedules use a tz-naive date index).
+    # Normalize to date-only boundaries, but ensure tz-awareness to match schedule index
     try:
-        start_day = pd.Timestamp(start_date.date())
+        start_day = pd.Timestamp(start_date.date(), tz=tzinfo)
     except Exception:
-        start_day = pd.Timestamp(pd.to_datetime(start_date).date())
+        start_day = pd.Timestamp(pd.to_datetime(start_date).date(), tz=tzinfo)
     try:
-        end_day_exclusive = pd.Timestamp(end_date.date())
+        end_day_exclusive = pd.Timestamp(end_date.date(), tz=tzinfo)
     except Exception:
-        end_day_exclusive = pd.Timestamp(pd.to_datetime(end_date).date())
+        end_day_exclusive = pd.Timestamp(pd.to_datetime(end_date).date(), tz=tzinfo)
 
     # Make end_date exclusive by moving it one day earlier.
     schedule_end_day = end_day_exclusive - pd.Timedelta(days=1)
@@ -229,6 +242,25 @@ def _get_year_schedule(year: int) -> pd.DataFrame:
 
     full_schedule = year_schedules[0] if len(year_schedules) == 1 else pd.concat(year_schedules, axis=0)
 
+    # Ensure the index is tz-aware and aligned with requested tz BEFORE slicing
+    # Calendars may return:
+    # - DatetimeIndex (tz-naive or tz-aware)
+    # - Index of Python datetimes (tz-naive or tz-aware)
+    idx = full_schedule.index
+    if isinstance(idx, pd.DatetimeIndex):
+        full_schedule.index = idx.tz_localize(tzinfo) if idx.tz is None else idx.tz_convert(tzinfo)
+    else:
+        # Likely an Index of Python datetime/date objects; can be mixed naive/aware
+        # Prefer parsing as naive and localizing (preserves calendar days)
+        # but fall back to utc=True if any tz-aware elements are present.
+        try:
+            tmp = pd.to_datetime(idx, errors="raise")
+            full_schedule.index = pd.DatetimeIndex(tmp).tz_localize(tzinfo)
+        except ValueError:
+            # Pandas requires utc=True when any tz-aware python datetimes are present
+            tmp = pd.to_datetime(idx, utc=True)
+            full_schedule.index = tmp.tz_convert(tzinfo)
+
     # Slice to the requested window (inclusive of schedule_end_day).
     days = full_schedule.loc[start_day:schedule_end_day].copy()
 

tests/backtest/backtest_performance_history.csv

@@ -7235,3 +7235,24 @@ timestamp,test_name,data_source,trading_days,execution_time_seconds,git_commit,l
 2026-01-27T19:02:10.083698,test_acceptance_ibkr_mes_futures_acceptance,unknown,,6.407,340ed4e7,4.4.41,,,,,Auto-tracked from test_acceptance_backtests_ci
 2026-01-27T21:02:44.774164,test_acceptance_backdoor_smartlimit,unknown,,210.067,8020e12e,4.4.41,,,,,Auto-tracked from test_acceptance_backtests_ci
 2026-01-27T21:06:26.716069,test_acceptance_spx_short_straddle,unknown,,213.94,8020e12e,4.4.41,,,,,Auto-tracked from test_acceptance_backtests_ci
+2026-02-01T21:17:26.655808,test_crypto_cash_regression_no_fees[price_map0],unknown,,3.302,395805d2,4.4.43,,,,,Auto-tracked from test_crypto_cash_regressions
+2026-02-01T21:17:29.922829,test_crypto_cash_regression_no_fees[price_map1],unknown,,3.15,395805d2,4.4.43,,,,,Auto-tracked from test_crypto_cash_regressions
+2026-02-01T21:17:35.109447,test_crypto_cash_regression_with_fees,unknown,,5.075,395805d2,4.4.43,,,,,Auto-tracked from test_crypto_cash_regressions
+2026-02-01T21:17:35.380215,test_databento_auth_failure_propagates,Databento,,0.155,395805d2,4.4.43,,,,,Auto-tracked from test_databento
+2026-02-01T21:17:47.312363,test_yahoo_finance_dividends,unknown,,11.801,395805d2,4.4.43,,,,,Auto-tracked from test_dividends
+2026-02-01T21:17:47.612178,test_stock_bracket,unknown,,0.18,395805d2,4.4.43,,,,,Auto-tracked from test_example_strategies
+2026-02-01T21:17:47.915589,test_stock_oco,unknown,,0.178,395805d2,4.4.43,,,,,Auto-tracked from test_example_strategies
+2026-02-01T21:17:48.343537,test_stock_buy_and_hold,unknown,,0.294,395805d2,4.4.43,,,,,Auto-tracked from test_example_strategies
+2026-02-01T21:17:49.334631,test_stock_diversified_leverage,unknown,,0.851,395805d2,4.4.43,,,,,Auto-tracked from test_example_strategies
+2026-02-01T21:17:49.668061,test_limit_and_trailing_stops,unknown,,0.18,395805d2,4.4.43,,,,,Auto-tracked from test_example_strategies
+2026-02-01T21:17:57.716977,test_trading_iteration_failure_raises_exception,unknown,,7.774,395805d2,4.4.43,,,,,Auto-tracked from test_failing_backtest
+2026-02-01T21:18:06.142554,test_backtest_classmethod_trading_iteration_failure,unknown,,8.135,395805d2,4.4.43,,,,,Auto-tracked from test_failing_backtest
+2026-02-01T21:18:11.924616,test_multileg_spread_backtest_cash_and_parent_fill,unknown,,4.979,395805d2,4.4.43,,,,,Auto-tracked from test_multileg_backtest
+2026-02-01T21:18:15.753524,test_multileg_parent_limit_order_fills_in_backtest,unknown,,3.664,395805d2,4.4.43,,,,,Auto-tracked from test_multileg_backtest
+2026-02-01T21:18:20.741212,test_pandas_datasource_with_daily_data_in_backtest,unknown,,2.966,395805d2,4.4.43,,,,,Auto-tracked from test_pandas_backtest
+2026-02-01T21:18:23.762486,test_bracket_orders_apply_entry_and_exit_fees,unknown,,2.858,395805d2,4.4.43,,,,,Auto-tracked from test_pandas_backtest
+2026-02-01T21:18:27.280865,test_bracket_positions_remain_bounded,unknown,,3.353,395805d2,4.4.43,,,,,Auto-tracked from test_pandas_backtest
+2026-02-01T21:18:31.517503,test_not_passing_trader_class_into_backtest_creates_generic_trader,unknown,,4.066,395805d2,4.4.43,,,,,Auto-tracked from test_passing_trader_into_backtest
+2026-02-01T21:18:35.813287,test_passing_trader_class_into_backtest_creates_trader_class,unknown,,4.127,395805d2,4.4.43,,,,,Auto-tracked from test_passing_trader_into_backtest
+2026-02-01T21:18:39.472260,test_s3_truncated_cache_forces_refetch,ThetaData,,0.105,395805d2,4.4.43,,,,,Auto-tracked from test_thetadata_resilience
+2026-02-01T21:18:43.490769,test_yahoo_last_price,Yahoo,,3.276,395805d2,4.4.43,,,,,Auto-tracked from test_yahoo

tests/test_helpers.py

@@ -4,6 +4,7 @@
 from zoneinfo import ZoneInfo
 import pytz
 import pytest
+import pandas as pd
 
 from lumibot import LUMIBOT_DEFAULT_TIMEZONE
 from lumibot.tools.helpers import (
@@ -372,3 +373,115 @@ def test_is_market_open_invalid_market():
     tz = pytz.timezone("US/Eastern")
     dtm = tz.localize(dt.datetime.combine(dt.date(2024, 1, 5), dt.time(10, 30)))
     assert is_market_open(dtm, "INVALID") is False
+
+
+@pytest.mark.parametrize(
+    "market, tzname",
+    [
+        ("NYSE", "America/New_York"),
+        # Using UTC also exhibits the same mismatch prior to the fix
+        ("NYSE", "UTC"),
+    ],
+)
+def test_get_trading_days_handles_tzaware_index_nyse(market, tzname):
+    """
+    Regression test for tz-aware index vs tz-naive slice bounds in get_trading_days.
+
+    Before the fix, this raises:
+        TypeError: Cannot compare tz-naive and tz-aware datetime-like objects
+
+    After the fix, it should return a non-empty schedule for the requested window.
+    """
+    tz = pytz.timezone(tzname)
+
+    # Intentionally pass timezone-aware datetimes with times that are not midnight
+    # so code paths normalize to date-only and then slice the tz-aware index.
+    start = tz.localize(dt.datetime(2025, 4, 14, 12, 0, 0))
+    end = tz.localize(dt.datetime(2025, 4, 20, 12, 0, 0))
+
+    sched = get_trading_days(market=market, start_date=start, end_date=end, tzinfo=tz)
+
+    # After fix: we should get a non-empty DataFrame with DatetimeIndex
+    assert isinstance(sched, pd.DataFrame)
+    assert not sched.empty
+    assert getattr(sched.index, "tz", None) is not None  # index should be tz-aware
+
+
+def test_get_trading_days_handles_tzaware_index_247():
+    """
+    Same regression test for the built-in "24/7" calendar.
+
+    Prior to the fix, this also triggers the tz-aware/naive slicing error.
+    After the fix, it should return daily sessions for the range.
+    """
+    tz = pytz.UTC
+    start = tz.localize(dt.datetime(2025, 1, 1, 8, 30, 0))
+    end = tz.localize(dt.datetime(2025, 1, 5, 17, 45, 0))
+
+    sched = get_trading_days(market="24/7", start_date=start, end_date=end, tzinfo=tz)
+
+    assert isinstance(sched, pd.DataFrame)
+    assert not sched.empty
+    # Expect 4 sessions: Jan 1, 2, 3, 4 (inclusive of start day; end is exclusive)
+    assert len(sched) == 4
+    assert getattr(sched.index, "tz", None) is not None
+
+
+def test_date_n_trading_days_from_date_no_tz_mismatch_nyse():
+    """
+    `date_n_trading_days_from_date` delegates to `get_trading_days`. Before the fix,
+    this call raises a tz-aware/naive mismatch TypeError. After the fix, it should
+    simply return a `datetime.date`.
+    """
+    tz = pytz.UTC
+    start_dt = tz.localize(dt.datetime(2025, 7, 1, 12, 0, 0))
+
+    result_back = date_n_trading_days_from_date(
+        n_days=5, start_datetime=start_dt, market="NYSE"
+    )
+    result_fwd = date_n_trading_days_from_date(
+        n_days=-5, start_datetime=start_dt, market="NYSE"
+    )
+
+    assert isinstance(result_back, dt.date)
+    assert isinstance(result_fwd, dt.date)
+
+
+def test_get_trading_days_on_tz_mismatch_then_fix(monkeypatch):
+    """
+    Construct a calendar whose schedule has a tz-aware DatetimeIndex (UTC),
+    while get_trading_days currently builds tz-naive slice bounds.
+    """
+    import pandas as pd
+    import pytz
+
+    class FakeCalendar:
+        def schedule(self, start_date, end_date, tz=None):
+            # Build a tz-aware DatetimeIndex to trigger the mismatch
+            idx = pd.date_range(
+                start=pd.Timestamp('2025-01-01', tz=pytz.UTC),
+                end=pd.Timestamp('2025-01-05', tz=pytz.UTC),
+                freq='D'
+            )
+            # Market open/close columns can be naive datetimes; they aren't
+            # used for the slicing that triggers the error.
+            opens = pd.date_range('2025-01-01 00:00:00', periods=len(idx), freq='D')
+            closes = pd.date_range('2025-01-01 23:59:00', periods=len(idx), freq='D')
+            df = pd.DataFrame({
+                'market_open': opens,
+                'market_close': closes,
+            }, index=idx)
+            return df
+
+    # Monkeypatch pandas_market_calendars.get_calendar used inside helpers
+    monkeypatch.setattr(helpers_module.mcal, 'get_calendar', lambda market: FakeCalendar())
+
+    tz = pytz.UTC
+    start = tz.localize(dt.datetime(2025, 1, 1, 12, 0, 0))
+    end = tz.localize(dt.datetime(2025, 1, 4, 12, 0, 0))
+
+    sched = get_trading_days(market="FAKE", start_date=start, end_date=end, tzinfo=tz)
+
+    assert isinstance(sched, pd.DataFrame)
+    assert not sched.empty
+    assert getattr(sched.index, 'tz', None) is not None