diff --git a/tests/calendars/test_trading_calendar.py b/tests/calendars/test_trading_calendar.py index 1cbb0b3a..9f6ef675 100644 --- a/tests/calendars/test_trading_calendar.py +++ b/tests/calendars/test_trading_calendar.py @@ -670,6 +670,14 @@ class ExchangeCalendarTestBase(object): found_open, found_close = \ self.calendar.open_and_close_for_session(session_label) + # Test that the methods for just session open and close produce the + # same values as the method for getting both. + alt_open = self.calendar.session_open(session_label) + self.assertEqual(alt_open, found_open) + + alt_close = self.calendar.session_close(session_label) + self.assertEqual(alt_close, found_close) + self.assertEqual(open_answer, found_open) self.assertEqual(close_answer, found_close) diff --git a/zipline/data/history_loader.py b/zipline/data/history_loader.py index df0b2789..e4503961 100644 --- a/zipline/data/history_loader.py +++ b/zipline/data/history_loader.py @@ -34,6 +34,7 @@ from zipline.lib.adjustment import Float64Multiply, Float64Add from zipline.utils.cache import ExpiringCache from zipline.utils.memoize import lazyval from zipline.utils.numpy_utils import float64_dtype +from zipline.utils.pandas_utils import find_in_sorted_index class HistoryCompatibleUSEquityAdjustmentReader(object): @@ -376,14 +377,10 @@ class HistoryLoader(with_metaclass(ABCMeta)): size = len(dts) asset_windows = {} needed_assets = [] + cal = self._calendar assets = self._asset_finder.retrieve_all(assets) - - try: - end_ix = self._calendar.searchsorted(end) - except KeyError: - raise KeyError("{0} not in calendar [{1}...{2}]".format( - end, self._calendar[0], self._calendar[-1])) + end_ix = find_in_sorted_index(cal, end) for asset in assets: try: @@ -401,15 +398,9 @@ class HistoryLoader(with_metaclass(ABCMeta)): asset_windows[asset] = window if needed_assets: - start = dts[0] - offset = 0 - try: - start_ix = self._calendar.searchsorted(start) - except KeyError: - raise KeyError("{0} not in calendar [{1}...{2}]".format( - start, self._calendar[0], self._calendar[-1])) - cal = self._calendar + start_ix = find_in_sorted_index(cal, dts[0]) + prefetch_end_ix = min(end_ix + self._prefetch_length, len(cal) - 1) prefetch_end = cal[prefetch_end_ix] prefetch_dts = cal[start_ix:prefetch_end_ix + 1] diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index 8dd09d21..8050162a 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -685,15 +685,25 @@ class WithEquityDailyBarData(WithTradingEnvironment): WithEquityMinuteBarData zipline.testing.create_daily_bar_data """ - EQUITY_DAILY_BAR_LOOKBACK_DAYS = 0 - EQUITY_DAILY_BAR_USE_FULL_CALENDAR = False EQUITY_DAILY_BAR_START_DATE = alias('START_DATE') EQUITY_DAILY_BAR_END_DATE = alias('END_DATE') EQUITY_DAILY_BAR_SOURCE_FROM_MINUTE = None + @classproperty + def EQUITY_DAILY_BAR_LOOKBACK_DAYS(cls): + # If we're sourcing from minute data, then we almost certainly want the + # minute bar calendar to be aligned with the daily bar calendar, so + # re-use the same lookback parameter. + if cls.EQUITY_DAILY_BAR_SOURCE_FROM_MINUTE: + return cls.EQUITY_MINUTE_BAR_LOOKBACK_DAYS + else: + return 0 + @classmethod def _make_equity_daily_bar_from_minute(cls): + assert issubclass(cls, WithEquityMinuteBarData), \ + "Can't source daily data from minute without minute data!" assets = cls.asset_finder.retrieve_all(cls.asset_finder.equities_sids) minute_data = dict(cls.make_equity_minute_bar_data()) for asset in assets: @@ -873,11 +883,9 @@ class WithEquityMinuteBarData(_WithMinuteBarDataBase): Methods ------- make_equity_minute_bar_data() -> iterable[(int, pd.DataFrame)] - A class method that returns a dict mapping sid to dataframe - which will be written to into the the format of the inherited - class which writes the minute bar data for use by a reader. - By default this creates some simple sythetic data with - :func:`~zipline.testing.create_minute_bar_data` + Classmethod producing an iterator of (sid, minute_data) pairs. + The default implementation invokes + zipline.testing.core.create_minute_bar_data. See Also -------- diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index c98419fc..05a7b104 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -643,6 +643,18 @@ class TradingCalendar(with_metaclass(ABCMeta)): return (o_and_c['market_open'].tz_localize('UTC'), o_and_c['market_close'].tz_localize('UTC')) + def session_open(self, session_label): + return self.schedule.loc[ + session_label, + 'market_open' + ].tz_localize('UTC') + + def session_close(self, session_label): + return self.schedule.loc[ + session_label, + 'market_close' + ].tz_localize('UTC') + @property def all_sessions(self): return self.schedule.index diff --git a/zipline/utils/pandas_utils.py b/zipline/utils/pandas_utils.py index 633960e1..ccac273a 100644 --- a/zipline/utils/pandas_utils.py +++ b/zipline/utils/pandas_utils.py @@ -91,6 +91,37 @@ def mask_between_time(dts, start, end, include_start=True, include_end=True): ) +def find_in_sorted_index(dts, dt): + """ + Find the index of ``dt`` in ``dts``. + + This function should be used instead of `dts.get_loc(dt)` if the index is + large enough that we don't want to initialize a hash table in ``dts``. In + particular, this should always be used on minutely trading calendars. + + Parameters + ---------- + dts : pd.DatetimeIndex + Index in which to look up ``dt``. **Must be sorted**. + dt : pd.Timestamp + ``dt`` to be looked up. + + Returns + ------- + ix : int + Integer index such that dts[ix] == dt. + + Raises + ------ + KeyError + If dt is not in ``dts``. + """ + ix = dts.searchsorted(dt) + if dts[ix] != dt: + raise LookupError("{dt} is not in {dts}".format(dt=dt, dts=dts)) + return ix + + def nearest_unequal_elements(dts, dt): """ Find values in ``dts`` closest but not equal to ``dt``.