diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index 77966c6a..1d0f6fdf 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABCMeta, abstractproperty +from lru import LRU from pandas.tseries.holiday import AbstractHolidayCalendar from six import with_metaclass @@ -95,6 +96,12 @@ class TradingCalendar(with_metaclass(ABCMeta)): dtype='datetime64[ns]', ) + # Simple cache to avoid recalculating the same minute -> session in + # "next" mode. Analysis of current zipline code paths show that + # `minute_to_session_label` is often called consecutively with the same + # inputs. + self._minute_to_session_label_cache = LRU(1) + self.market_opens_nanos = self.schedule.market_open.values.\ astype(np.int64) @@ -690,11 +697,19 @@ class TradingCalendar(with_metaclass(ABCMeta)): pd.Timestamp (midnight UTC) The label of the containing session. """ + if direction == "next": + try: + return self._minute_to_session_label_cache[dt] + except KeyError: + pass idx = searchsorted(self.market_closes_nanos, dt) current_or_next_session = self.schedule.index[idx] + self._minute_to_session_label_cache[dt] = current_or_next_session - if direction == "previous": + if direction == "next": + return current_or_next_session + elif direction == "previous": if not is_open(self.market_opens_nanos, self.market_closes_nanos, dt): # if the exchange is closed, use the previous session @@ -704,7 +719,7 @@ class TradingCalendar(with_metaclass(ABCMeta)): dt): # if the exchange is closed, blow up raise ValueError("The given dt is not an exchange minute!") - elif direction != "next": + else: # invalid direction raise ValueError("Invalid direction parameter: " "{0}".format(direction))