mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-27 19:30:28 +08:00
ENH: Add new parameter to schedule_function that accepts a trading
calendar.
This commit is contained in:
Jean Bredeche
@@ -16,14 +16,12 @@ import datetime
|
||||
from inspect import isabstract
|
||||
import random
|
||||
from unittest import TestCase
|
||||
from datetime import timedelta
|
||||
|
||||
from nose_parameterized import parameterized
|
||||
import pandas as pd
|
||||
from six import iteritems
|
||||
from six.moves import range, map
|
||||
|
||||
from zipline.testing import subtest, parameter_space
|
||||
import zipline.utils.events
|
||||
from zipline.utils.calendars import get_calendar
|
||||
from zipline.utils.events import (
|
||||
@@ -182,10 +180,10 @@ class TestEventRule(TestCase):
|
||||
super(Always, Always()).should_trigger('a')
|
||||
|
||||
|
||||
def minutes_for_days(ordered_days=False):
|
||||
def minutes_for_days(cal, ordered_days=False):
|
||||
"""
|
||||
500 randomly selected days.
|
||||
This is used to make sure our test coverage is unbaised towards any rules.
|
||||
This is used to make sure our test coverage is unbiased towards any rules.
|
||||
We use a random sample because testing on all the trading days took
|
||||
around 180 seconds on my laptop, which is far too much for normal unit
|
||||
testing.
|
||||
@@ -198,7 +196,6 @@ def minutes_for_days(ordered_days=False):
|
||||
Iterating over this yields a single day, iterating over the day yields
|
||||
the minutes for that day.
|
||||
"""
|
||||
cal = get_calendar('NYSE')
|
||||
random.seed('deterministic')
|
||||
if ordered_days:
|
||||
# Get a list of 500 trading days, in order. As a performance
|
||||
@@ -216,11 +213,13 @@ def minutes_for_days(ordered_days=False):
|
||||
def session_picker(day):
|
||||
return random.choice(cal.all_sessions[:-1])
|
||||
|
||||
return ((cal.minutes_for_session(session_picker(cnt)),)
|
||||
for cnt in range(500))
|
||||
return [cal.minutes_for_session(session_picker(cnt))
|
||||
for cnt in range(500)]
|
||||
|
||||
|
||||
class RuleTestCase(TestCase):
|
||||
class RuleTestCase(object):
|
||||
CALENDAR_STRING = "foo"
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# On the AfterOpen and BeforeClose tests, we want ensure that the
|
||||
@@ -234,7 +233,7 @@ class RuleTestCase(TestCase):
|
||||
cls.after_open = AfterOpen(hours=1, minutes=5)
|
||||
cls.class_ = None # Mark that this is the base class.
|
||||
|
||||
cal = get_calendar('NYSE')
|
||||
cal = get_calendar(cls.CALENDAR_STRING)
|
||||
cls.before_close.cal = cal
|
||||
cls.after_open.cal = cal
|
||||
|
||||
@@ -266,289 +265,179 @@ class RuleTestCase(TestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestStatelessRules(RuleTestCase):
|
||||
class StatelessRulesTests(RuleTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super(TestStatelessRules, cls).setUpClass()
|
||||
super(StatelessRulesTests, cls).setUpClass()
|
||||
|
||||
cls.class_ = StatelessRule
|
||||
|
||||
cls.nyse_cal = get_calendar('NYSE')
|
||||
cls.cal = get_calendar(cls.CALENDAR_STRING)
|
||||
|
||||
# First day of 09/2014 is closed whereas that for 10/2014 is open
|
||||
cls.sept_sessions = cls.nyse_cal.sessions_in_range(
|
||||
cls.sept_sessions = cls.cal.sessions_in_range(
|
||||
pd.Timestamp('2014-09-01', tz='UTC'),
|
||||
pd.Timestamp('2014-09-30', tz='UTC'),
|
||||
)
|
||||
cls.oct_sessions = cls.nyse_cal.sessions_in_range(
|
||||
cls.oct_sessions = cls.cal.sessions_in_range(
|
||||
pd.Timestamp('2014-10-01', tz='UTC'),
|
||||
pd.Timestamp('2014-10-31', tz='UTC'),
|
||||
)
|
||||
|
||||
cls.sept_week = cls.nyse_cal.minutes_for_sessions_in_range(
|
||||
cls.sept_week = cls.cal.minutes_for_sessions_in_range(
|
||||
pd.Timestamp("2014-09-22", tz='UTC'),
|
||||
pd.Timestamp("2014-09-26", tz='UTC')
|
||||
)
|
||||
|
||||
@subtest(minutes_for_days(), 'ms')
|
||||
def test_Always(self, ms):
|
||||
cls.HALF_SESSION = None
|
||||
cls.FULL_SESSION = None
|
||||
|
||||
def test_Always(self):
|
||||
should_trigger = Always().should_trigger
|
||||
self.assertTrue(all(map(should_trigger, ms)))
|
||||
for session_minutes in minutes_for_days(self.cal):
|
||||
self.assertTrue(all(map(should_trigger, session_minutes)))
|
||||
|
||||
@subtest(minutes_for_days(), 'ms')
|
||||
def test_Never(self, ms):
|
||||
def test_Never(self):
|
||||
should_trigger = Never().should_trigger
|
||||
self.assertFalse(any(map(should_trigger, ms)))
|
||||
for session_minutes in minutes_for_days(self.cal):
|
||||
self.assertFalse(any(map(should_trigger, session_minutes)))
|
||||
|
||||
@subtest(minutes_for_days(ordered_days=True), 'ms')
|
||||
def test_AfterOpen(self, ms):
|
||||
def test_AfterOpen(self):
|
||||
minute_groups = minutes_for_days(self.cal, ordered_days=True)
|
||||
should_trigger = self.after_open.should_trigger
|
||||
for i, m in enumerate(ms):
|
||||
# Should only trigger at the 64th minute
|
||||
if i != 64:
|
||||
self.assertFalse(should_trigger(m))
|
||||
else:
|
||||
self.assertTrue(should_trigger(m))
|
||||
for session_minutes in minute_groups:
|
||||
for i, minute in enumerate(session_minutes):
|
||||
# Should only trigger at the 64th minute
|
||||
if i != 64:
|
||||
self.assertFalse(should_trigger(minute))
|
||||
else:
|
||||
self.assertTrue(should_trigger(minute))
|
||||
|
||||
@subtest(minutes_for_days(ordered_days=True), 'ms')
|
||||
def test_BeforeClose(self, ms):
|
||||
ms = list(ms)
|
||||
def test_BeforeClose(self):
|
||||
minute_groups = minutes_for_days(self.cal, ordered_days=True)
|
||||
should_trigger = self.before_close.should_trigger
|
||||
for m in ms:
|
||||
# Should only trigger at the 65th-to-last minute
|
||||
if m != ms[-66]:
|
||||
self.assertFalse(should_trigger(m))
|
||||
else:
|
||||
self.assertTrue(should_trigger(m))
|
||||
for minute_group in minute_groups:
|
||||
for minute in minute_group:
|
||||
# Should only trigger at the 65th-to-last minute
|
||||
if minute != minute_group[-66]:
|
||||
self.assertFalse(should_trigger(minute))
|
||||
else:
|
||||
self.assertTrue(should_trigger(minute))
|
||||
|
||||
def test_NotHalfDay(self):
|
||||
rule = NotHalfDay()
|
||||
rule.cal = self.nyse_cal
|
||||
rule.cal = self.cal
|
||||
|
||||
half_day_period = pd.Timestamp("2014-07-03", tz='UTC')
|
||||
full_day_period = pd.Timestamp("2014-09-24", tz='UTC')
|
||||
if self.HALF_SESSION:
|
||||
for minute in self.cal.minutes_for_session(self.HALF_SESSION):
|
||||
self.assertFalse(rule.should_trigger(minute))
|
||||
|
||||
for minute in self.nyse_cal.minutes_for_session(half_day_period):
|
||||
self.assertFalse(rule.should_trigger(minute))
|
||||
|
||||
for minute in self.nyse_cal.minutes_for_session(full_day_period):
|
||||
self.assertTrue(rule.should_trigger(minute))
|
||||
if self.FULL_SESSION:
|
||||
for minute in self.cal.minutes_for_session(self.FULL_SESSION):
|
||||
self.assertTrue(rule.should_trigger(minute))
|
||||
|
||||
def test_NthTradingDayOfWeek_day_zero(self):
|
||||
"""
|
||||
Test that we don't blow up when trying to call week_start's
|
||||
should_trigger on the first day of a trading environment.
|
||||
"""
|
||||
cal = get_calendar('NYSE')
|
||||
rule = NthTradingDayOfWeek(0)
|
||||
rule.cal = cal
|
||||
first_open = self.nyse_cal.open_and_close_for_session(
|
||||
self.nyse_cal.all_sessions[0]
|
||||
rule.cal = self.cal
|
||||
first_open = self.cal.open_and_close_for_session(
|
||||
self.cal.all_sessions[0]
|
||||
)
|
||||
self.assertTrue(first_open)
|
||||
|
||||
@subtest(param_range(MAX_WEEK_RANGE), 'n')
|
||||
def test_NthTradingDayOfWeek(self, n):
|
||||
cal = get_calendar('NYSE')
|
||||
rule = NthTradingDayOfWeek(n)
|
||||
rule.cal = cal
|
||||
should_trigger = rule.should_trigger
|
||||
prev_period = self.nyse_cal.minute_to_session_label(self.sept_week[0])
|
||||
n_tdays = 0
|
||||
for minute in self.sept_week:
|
||||
period = self.nyse_cal.minute_to_session_label(
|
||||
minute, direction="none"
|
||||
)
|
||||
def test_NthTradingDayOfWeek(self):
|
||||
for n in range(MAX_WEEK_RANGE):
|
||||
rule = NthTradingDayOfWeek(n)
|
||||
rule.cal = self.cal
|
||||
should_trigger = rule.should_trigger
|
||||
prev_period = self.cal.minute_to_session_label(self.sept_week[0])
|
||||
n_tdays = 0
|
||||
for minute in self.sept_week:
|
||||
period = self.cal.minute_to_session_label(minute)
|
||||
|
||||
if prev_period < period:
|
||||
n_tdays += 1
|
||||
prev_period = period
|
||||
|
||||
if should_trigger(minute):
|
||||
self.assertEqual(n_tdays, n)
|
||||
else:
|
||||
self.assertNotEqual(n_tdays, n)
|
||||
|
||||
@subtest(param_range(MAX_WEEK_RANGE), 'n')
|
||||
def test_NDaysBeforeLastTradingDayOfWeek(self, n):
|
||||
cal = get_calendar('NYSE')
|
||||
rule = NDaysBeforeLastTradingDayOfWeek(n)
|
||||
rule.cal = cal
|
||||
should_trigger = rule.should_trigger
|
||||
for minute in self.sept_week:
|
||||
if should_trigger(minute):
|
||||
n_tdays = 0
|
||||
session = self.nyse_cal.minute_to_session_label(
|
||||
minute,
|
||||
direction="none"
|
||||
)
|
||||
next_session = self.nyse_cal.next_session_label(session)
|
||||
while next_session.dayofweek > session.dayofweek:
|
||||
session = next_session
|
||||
next_session = self.nyse_cal.next_session_label(session)
|
||||
if prev_period < period:
|
||||
n_tdays += 1
|
||||
prev_period = period
|
||||
|
||||
self.assertEqual(n_tdays, n)
|
||||
|
||||
@parameter_space(
|
||||
rule_offset=(0, 1, 2, 3, 4),
|
||||
start_offset=(0, 1, 2, 3, 4),
|
||||
type=('week_start', 'week_end')
|
||||
)
|
||||
def test_edge_cases_for_TradingDayOfWeek(self,
|
||||
rule_offset,
|
||||
start_offset,
|
||||
type):
|
||||
"""
|
||||
Test that we account for midweek holidays. Monday 01/20 is a holiday.
|
||||
Ensure that the trigger date for that week is adjusted
|
||||
appropriately, or thrown out if not enough trading days. Also, test
|
||||
that if we start the simulation on a day where we miss the trigger
|
||||
for that week, that the trigger is recalculated for next week.
|
||||
"""
|
||||
|
||||
sim_start = pd.Timestamp('2014-01-06', tz='UTC') + \
|
||||
timedelta(days=start_offset)
|
||||
|
||||
delta = timedelta(days=start_offset)
|
||||
|
||||
jan_minutes = self.nyse_cal.minutes_for_sessions_in_range(
|
||||
pd.Timestamp("2014-01-06", tz='UTC') + delta,
|
||||
pd.Timestamp("2014-01-31", tz='UTC')
|
||||
)
|
||||
|
||||
if type == 'week_start':
|
||||
rule = NthTradingDayOfWeek
|
||||
# Expect to trigger on the first trading day of the week, plus the
|
||||
# offset
|
||||
trigger_periods = [
|
||||
pd.Timestamp('2014-01-06', tz='UTC'),
|
||||
pd.Timestamp('2014-01-13', tz='UTC'),
|
||||
pd.Timestamp('2014-01-21', tz='UTC'),
|
||||
pd.Timestamp('2014-01-27', tz='UTC'),
|
||||
]
|
||||
trigger_periods = \
|
||||
[x + timedelta(days=rule_offset) for x in trigger_periods]
|
||||
else:
|
||||
rule = NDaysBeforeLastTradingDayOfWeek
|
||||
# Expect to trigger on the last trading day of the week, minus the
|
||||
# offset
|
||||
trigger_periods = [
|
||||
pd.Timestamp('2014-01-10', tz='UTC'),
|
||||
pd.Timestamp('2014-01-17', tz='UTC'),
|
||||
pd.Timestamp('2014-01-24', tz='UTC'),
|
||||
pd.Timestamp('2014-01-31', tz='UTC'),
|
||||
]
|
||||
trigger_periods = \
|
||||
[x - timedelta(days=rule_offset) for x in trigger_periods]
|
||||
|
||||
rule.cal = self.nyse_cal
|
||||
should_trigger = rule(rule_offset).should_trigger
|
||||
|
||||
# If offset is 4, there is not enough trading days in the short week,
|
||||
# and so it should not trigger
|
||||
if rule_offset == 4:
|
||||
del trigger_periods[2]
|
||||
|
||||
# Filter out trigger dates that happen before the simulation starts
|
||||
trigger_periods = [x for x in trigger_periods if x >= sim_start]
|
||||
|
||||
# Get all the minutes on the trigger dates
|
||||
trigger_minutes = self.nyse_cal.minutes_for_session(trigger_periods[0])
|
||||
for period in trigger_periods[1:]:
|
||||
trigger_minutes += self.nyse_cal.minutes_for_session(period)
|
||||
|
||||
expected_n_triggered = len(trigger_minutes)
|
||||
trigger_minutes_iter = iter(trigger_minutes)
|
||||
|
||||
n_triggered = 0
|
||||
for m in jan_minutes:
|
||||
if should_trigger(m):
|
||||
self.assertEqual(m, next(trigger_minutes_iter))
|
||||
n_triggered += 1
|
||||
|
||||
self.assertEqual(n_triggered, expected_n_triggered)
|
||||
|
||||
@parameterized.expand([('week_start',), ('week_end',)])
|
||||
def test_week_and_time_composed_rule(self, type):
|
||||
week_rule = NthTradingDayOfWeek(0) if type == 'week_start' else \
|
||||
NDaysBeforeLastTradingDayOfWeek(4)
|
||||
time_rule = AfterOpen(minutes=60)
|
||||
|
||||
week_rule.cal = self.nyse_cal
|
||||
time_rule.cal = self.nyse_cal
|
||||
|
||||
composed_rule = week_rule & time_rule
|
||||
|
||||
should_trigger = composed_rule.should_trigger
|
||||
|
||||
week_minutes = self.nyse_cal.minutes_for_sessions_in_range(
|
||||
pd.Timestamp("2014-01-06", tz='UTC'),
|
||||
pd.Timestamp("2014-01-10", tz='UTC')
|
||||
)
|
||||
|
||||
dt = pd.Timestamp('2014-01-06 14:30:00', tz='UTC')
|
||||
trigger_day_offset = 0
|
||||
trigger_minute_offset = 60
|
||||
n_triggered = 0
|
||||
|
||||
for m in week_minutes:
|
||||
if should_trigger(m):
|
||||
self.assertEqual(m, dt + timedelta(days=trigger_day_offset) +
|
||||
timedelta(minutes=trigger_minute_offset))
|
||||
n_triggered += 1
|
||||
|
||||
self.assertEqual(n_triggered, 1)
|
||||
|
||||
@subtest(param_range(MAX_MONTH_RANGE), 'n')
|
||||
def test_NthTradingDayOfMonth(self, n):
|
||||
cal = get_calendar('NYSE')
|
||||
rule = NthTradingDayOfMonth(n)
|
||||
rule.cal = cal
|
||||
should_trigger = rule.should_trigger
|
||||
for sessions_list in (self.sept_sessions, self.oct_sessions):
|
||||
for n_tdays, session in enumerate(sessions_list):
|
||||
for m in self.nyse_cal.minutes_for_session(session):
|
||||
if should_trigger(m):
|
||||
self.assertEqual(n_tdays, n)
|
||||
else:
|
||||
self.assertNotEqual(n_tdays, n)
|
||||
|
||||
@subtest(param_range(MAX_MONTH_RANGE), 'n')
|
||||
def test_NDaysBeforeLastTradingDayOfMonth(self, n):
|
||||
cal = get_calendar('NYSE')
|
||||
rule = NDaysBeforeLastTradingDayOfMonth(n)
|
||||
rule.cal = cal
|
||||
should_trigger = rule.should_trigger
|
||||
for n_days_before, session in enumerate(reversed(self.oct_sessions)):
|
||||
for m in self.nyse_cal.minutes_for_session(session):
|
||||
if should_trigger(m):
|
||||
self.assertEqual(n_days_before, n)
|
||||
if should_trigger(minute):
|
||||
self.assertEqual(n_tdays, n)
|
||||
else:
|
||||
self.assertNotEqual(n_days_before, n)
|
||||
self.assertNotEqual(n_tdays, n)
|
||||
|
||||
@subtest(minutes_for_days(), 'ms')
|
||||
def test_ComposedRule(self, ms):
|
||||
def test_NDaysBeforeLastTradingDayOfWeek(self):
|
||||
for n in range(MAX_WEEK_RANGE):
|
||||
rule = NDaysBeforeLastTradingDayOfWeek(n)
|
||||
rule.cal = self.cal
|
||||
should_trigger = rule.should_trigger
|
||||
for minute in self.sept_week:
|
||||
if should_trigger(minute):
|
||||
n_tdays = 0
|
||||
session = self.cal.minute_to_session_label(
|
||||
minute,
|
||||
direction="none"
|
||||
)
|
||||
next_session = self.cal.next_session_label(session)
|
||||
while next_session.dayofweek > session.dayofweek:
|
||||
session = next_session
|
||||
next_session = self.cal.next_session_label(session)
|
||||
n_tdays += 1
|
||||
|
||||
self.assertEqual(n_tdays, n)
|
||||
|
||||
def test_NthTradingDayOfMonth(self):
|
||||
for n in range(MAX_MONTH_RANGE):
|
||||
rule = NthTradingDayOfMonth(n)
|
||||
rule.cal = self.cal
|
||||
should_trigger = rule.should_trigger
|
||||
for sessions_list in (self.sept_sessions, self.oct_sessions):
|
||||
for n_tdays, session in enumerate(sessions_list):
|
||||
# just check the first 10 minutes of each session
|
||||
for m in self.cal.minutes_for_session(session)[0:10]:
|
||||
if should_trigger(m):
|
||||
self.assertEqual(n_tdays, n)
|
||||
else:
|
||||
self.assertNotEqual(n_tdays, n)
|
||||
|
||||
def test_NDaysBeforeLastTradingDayOfMonth(self):
|
||||
for n in range(MAX_MONTH_RANGE):
|
||||
rule = NDaysBeforeLastTradingDayOfMonth(n)
|
||||
rule.cal = self.cal
|
||||
should_trigger = rule.should_trigger
|
||||
sessions = reversed(self.oct_sessions)
|
||||
for n_days_before, session in enumerate(sessions):
|
||||
for m in self.cal.minutes_for_session(session)[0:10]:
|
||||
if should_trigger(m):
|
||||
self.assertEqual(n_days_before, n)
|
||||
else:
|
||||
self.assertNotEqual(n_days_before, n)
|
||||
|
||||
def test_ComposedRule(self):
|
||||
minute_groups = minutes_for_days(self.cal)
|
||||
rule1 = Always()
|
||||
rule2 = Never()
|
||||
|
||||
composed = rule1 & rule2
|
||||
should_trigger = composed.should_trigger
|
||||
self.assertIsInstance(composed, ComposedRule)
|
||||
self.assertIs(composed.first, rule1)
|
||||
self.assertIs(composed.second, rule2)
|
||||
self.assertFalse(any(map(should_trigger, ms)))
|
||||
for minute in minute_groups:
|
||||
composed = rule1 & rule2
|
||||
should_trigger = composed.should_trigger
|
||||
self.assertIsInstance(composed, ComposedRule)
|
||||
self.assertIs(composed.first, rule1)
|
||||
self.assertIs(composed.second, rule2)
|
||||
self.assertFalse(any(map(should_trigger, minute)))
|
||||
|
||||
|
||||
class TestStatefulRules(RuleTestCase):
|
||||
class StatefulRulesTests(RuleTestCase):
|
||||
CALENDAR_STRING = "NYSE"
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super(TestStatefulRules, cls).setUpClass()
|
||||
super(StatefulRulesTests, cls).setUpClass()
|
||||
|
||||
cls.class_ = StatefulRule
|
||||
cls.cal = get_calendar(cls.CALENDAR_STRING)
|
||||
|
||||
@subtest(minutes_for_days(), 'ms')
|
||||
def test_OncePerDay(self, ms):
|
||||
def test_OncePerDay(self):
|
||||
class RuleCounter(StatefulRule):
|
||||
"""
|
||||
A rule that counts the number of times another rule triggers
|
||||
@@ -562,8 +451,10 @@ class TestStatefulRules(RuleTestCase):
|
||||
self.count += 1
|
||||
return st
|
||||
|
||||
rule = RuleCounter(OncePerDay())
|
||||
for m in ms:
|
||||
rule.should_trigger(m)
|
||||
for minute_group in minutes_for_days(self.cal):
|
||||
rule = RuleCounter(OncePerDay())
|
||||
|
||||
self.assertEqual(rule.count, 1)
|
||||
for minute in minute_group:
|
||||
rule.should_trigger(minute)
|
||||
|
||||
self.assertEqual(rule.count, 1)
|
||||
@@ -0,0 +1,29 @@
|
||||
#
|
||||
# Copyright 2016 Quantopian, Inc.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from unittest import TestCase
|
||||
import pandas as pd
|
||||
|
||||
from test_events import StatefulRulesTests, StatelessRulesTests
|
||||
|
||||
|
||||
class TestStatelessRulesCME(StatelessRulesTests, TestCase):
|
||||
CALENDAR_STRING = "CME"
|
||||
|
||||
HALF_SESSION = pd.Timestamp("2014-07-04", tz='UTC')
|
||||
FULL_SESSION = pd.Timestamp("2014-09-24", tz='UTC')
|
||||
|
||||
|
||||
class TestStatefulRulesCME(StatefulRulesTests, TestCase):
|
||||
CALENDAR_STRING = "CME"
|
||||
@@ -0,0 +1,145 @@
|
||||
#
|
||||
# Copyright 2016 Quantopian, Inc.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from unittest import TestCase
|
||||
from datetime import timedelta
|
||||
import pandas as pd
|
||||
from nose_parameterized import parameterized
|
||||
|
||||
from zipline.testing import parameter_space
|
||||
from zipline.utils.events import NDaysBeforeLastTradingDayOfWeek, AfterOpen
|
||||
from zipline.utils.events import NthTradingDayOfWeek
|
||||
|
||||
from test_events import StatelessRulesTests, StatefulRulesTests
|
||||
|
||||
|
||||
class TestStatelessRulesNYSE(StatelessRulesTests, TestCase):
|
||||
CALENDAR_STRING = "NYSE"
|
||||
|
||||
HALF_SESSION = pd.Timestamp("2014-07-03", tz='UTC')
|
||||
FULL_SESSION = pd.Timestamp("2014-09-24", tz='UTC')
|
||||
|
||||
@parameter_space(
|
||||
rule_offset=(0, 1, 2, 3, 4),
|
||||
start_offset=(0, 1, 2, 3, 4),
|
||||
type=('week_start', 'week_end')
|
||||
)
|
||||
def test_edge_cases_for_TradingDayOfWeek(self,
|
||||
rule_offset,
|
||||
start_offset,
|
||||
type):
|
||||
"""
|
||||
Test that we account for midweek holidays. Monday 01/20 is a holiday.
|
||||
Ensure that the trigger date for that week is adjusted
|
||||
appropriately, or thrown out if not enough trading days. Also, test
|
||||
that if we start the simulation on a day where we miss the trigger
|
||||
for that week, that the trigger is recalculated for next week.
|
||||
"""
|
||||
|
||||
sim_start = pd.Timestamp('2014-01-06', tz='UTC') + \
|
||||
timedelta(days=start_offset)
|
||||
|
||||
delta = timedelta(days=start_offset)
|
||||
|
||||
jan_minutes = self.cal.minutes_for_sessions_in_range(
|
||||
pd.Timestamp("2014-01-06", tz='UTC') + delta,
|
||||
pd.Timestamp("2014-01-31", tz='UTC')
|
||||
)
|
||||
|
||||
if type == 'week_start':
|
||||
rule = NthTradingDayOfWeek
|
||||
# Expect to trigger on the first trading day of the week, plus the
|
||||
# offset
|
||||
trigger_periods = [
|
||||
pd.Timestamp('2014-01-06', tz='UTC'),
|
||||
pd.Timestamp('2014-01-13', tz='UTC'),
|
||||
pd.Timestamp('2014-01-21', tz='UTC'),
|
||||
pd.Timestamp('2014-01-27', tz='UTC'),
|
||||
]
|
||||
trigger_periods = \
|
||||
[x + timedelta(days=rule_offset) for x in trigger_periods]
|
||||
else:
|
||||
rule = NDaysBeforeLastTradingDayOfWeek
|
||||
# Expect to trigger on the last trading day of the week, minus the
|
||||
# offset
|
||||
trigger_periods = [
|
||||
pd.Timestamp('2014-01-10', tz='UTC'),
|
||||
pd.Timestamp('2014-01-17', tz='UTC'),
|
||||
pd.Timestamp('2014-01-24', tz='UTC'),
|
||||
pd.Timestamp('2014-01-31', tz='UTC'),
|
||||
]
|
||||
trigger_periods = \
|
||||
[x - timedelta(days=rule_offset) for x in trigger_periods]
|
||||
|
||||
rule.cal = self.cal
|
||||
should_trigger = rule(rule_offset).should_trigger
|
||||
|
||||
# If offset is 4, there is not enough trading days in the short week,
|
||||
# and so it should not trigger
|
||||
if rule_offset == 4:
|
||||
del trigger_periods[2]
|
||||
|
||||
# Filter out trigger dates that happen before the simulation starts
|
||||
trigger_periods = [x for x in trigger_periods if x >= sim_start]
|
||||
|
||||
# Get all the minutes on the trigger dates
|
||||
trigger_minutes = self.cal.minutes_for_session(trigger_periods[0])
|
||||
for period in trigger_periods[1:]:
|
||||
trigger_minutes += self.cal.minutes_for_session(period)
|
||||
|
||||
expected_n_triggered = len(trigger_minutes)
|
||||
trigger_minutes_iter = iter(trigger_minutes)
|
||||
|
||||
n_triggered = 0
|
||||
for m in jan_minutes:
|
||||
if should_trigger(m):
|
||||
self.assertEqual(m, next(trigger_minutes_iter))
|
||||
n_triggered += 1
|
||||
|
||||
self.assertEqual(n_triggered, expected_n_triggered)
|
||||
|
||||
@parameterized.expand([('week_start',), ('week_end',)])
|
||||
def test_week_and_time_composed_rule(self, type):
|
||||
week_rule = NthTradingDayOfWeek(0) if type == 'week_start' else \
|
||||
NDaysBeforeLastTradingDayOfWeek(4)
|
||||
time_rule = AfterOpen(minutes=60)
|
||||
|
||||
week_rule.cal = self.cal
|
||||
time_rule.cal = self.cal
|
||||
|
||||
composed_rule = week_rule & time_rule
|
||||
|
||||
should_trigger = composed_rule.should_trigger
|
||||
|
||||
week_minutes = self.cal.minutes_for_sessions_in_range(
|
||||
pd.Timestamp("2014-01-06", tz='UTC'),
|
||||
pd.Timestamp("2014-01-10", tz='UTC')
|
||||
)
|
||||
|
||||
dt = pd.Timestamp('2014-01-06 14:30:00', tz='UTC')
|
||||
trigger_day_offset = 0
|
||||
trigger_minute_offset = 60
|
||||
n_triggered = 0
|
||||
|
||||
for m in week_minutes:
|
||||
if should_trigger(m):
|
||||
self.assertEqual(m, dt + timedelta(days=trigger_day_offset) +
|
||||
timedelta(minutes=trigger_minute_offset))
|
||||
n_triggered += 1
|
||||
|
||||
self.assertEqual(n_triggered, 1)
|
||||
|
||||
|
||||
class TestStatefulRulesNYSE(StatefulRulesTests, TestCase):
|
||||
CALENDAR_STRING = "NYSE"
|
||||
@@ -422,6 +422,61 @@ def handle_data(context, data):
|
||||
env=self.env)
|
||||
algo.run(self.data_portal)
|
||||
|
||||
def test_schedule_function_custom_cal(self):
|
||||
# run a simulation on the CME cal, and schedule a function
|
||||
# using the NYSE cal
|
||||
algotext = """
|
||||
from zipline.api import schedule_function, get_datetime, time_rules, date_rules
|
||||
from zipline.utils.calendars import get_calendar
|
||||
|
||||
def initialize(context):
|
||||
schedule_function(
|
||||
func=log_nyse_open,
|
||||
date_rule=date_rules.every_day(),
|
||||
time_rule=time_rules.market_open(),
|
||||
calendar=get_calendar("NYSE")
|
||||
)
|
||||
|
||||
schedule_function(
|
||||
func=log_nyse_close,
|
||||
date_rule=date_rules.every_day(),
|
||||
time_rule=time_rules.market_close(),
|
||||
calendar=get_calendar("NYSE")
|
||||
)
|
||||
|
||||
context.nyse_opens = []
|
||||
context.nyse_closes = []
|
||||
|
||||
def log_nyse_open(context, data):
|
||||
context.nyse_opens.append(get_datetime())
|
||||
|
||||
def log_nyse_close(context, data):
|
||||
context.nyse_closes.append(get_datetime())
|
||||
"""
|
||||
|
||||
algo = TradingAlgorithm(
|
||||
script=algotext,
|
||||
sim_params=self.sim_params,
|
||||
env=self.env,
|
||||
trading_calendar=get_calendar("CME")
|
||||
)
|
||||
|
||||
algo.run(self.data_portal)
|
||||
|
||||
nyse = get_calendar("NYSE")
|
||||
|
||||
for minute in algo.nyse_opens:
|
||||
# each minute should be a nyse session open
|
||||
session_label = nyse.minute_to_session_label(minute)
|
||||
session_open = nyse.open_and_close_for_session(session_label)[0]
|
||||
self.assertEqual(session_open, minute)
|
||||
|
||||
for minute in algo.nyse_closes:
|
||||
# each minute should be a minute before a nyse session close
|
||||
session_label = nyse.minute_to_session_label(minute)
|
||||
session_close = nyse.open_and_close_for_session(session_label)[1]
|
||||
self.assertEqual(session_close - timedelta(minutes=1), minute)
|
||||
|
||||
def test_schedule_function(self):
|
||||
us_eastern = pytz.timezone('US/Eastern')
|
||||
|
||||
|
||||
@@ -1044,7 +1044,8 @@ class TradingAlgorithm(object):
|
||||
func,
|
||||
date_rule=None,
|
||||
time_rule=None,
|
||||
half_days=True):
|
||||
half_days=True,
|
||||
calendar=None):
|
||||
"""Schedules a function to be called according to some timed rules.
|
||||
|
||||
Parameters
|
||||
@@ -1082,8 +1083,7 @@ class TradingAlgorithm(object):
|
||||
# Check the type of the algorithm's schedule before pulling calendar
|
||||
# Note that the ExchangeTradingSchedule is currently the only
|
||||
# TradingSchedule class, so this is unlikely to be hit
|
||||
# TODO The calendar should be a required arg for schedule_function
|
||||
cal = self.trading_calendar
|
||||
cal = calendar or self.trading_calendar
|
||||
|
||||
self.add_event(
|
||||
make_eventrule(date_rule, time_rule, cal, half_days),
|
||||
|
||||
@@ -321,7 +321,7 @@ class AfterOpen(StatelessRule):
|
||||
# given a dt, find that day's open and period end (open + offset)
|
||||
self._period_start, self._period_close = \
|
||||
self.cal.open_and_close_for_session(
|
||||
self.cal.minute_to_session_label(dt, direction="none")
|
||||
self.cal.minute_to_session_label(dt)
|
||||
)
|
||||
self._period_end = self._period_start + self.offset - self._one_minute
|
||||
|
||||
@@ -396,7 +396,7 @@ class NotHalfDay(StatelessRule):
|
||||
A rule that only triggers when it is not a half day.
|
||||
"""
|
||||
def should_trigger(self, dt):
|
||||
return self.cal.minute_to_session_label(dt, direction="none") \
|
||||
return self.cal.minute_to_session_label(dt) \
|
||||
not in self.cal.early_closes
|
||||
|
||||
|
||||
@@ -416,7 +416,7 @@ class TradingDayOfWeekRule(six.with_metaclass(ABCMeta, StatelessRule)):
|
||||
|
||||
def should_trigger(self, dt):
|
||||
# is this market minute's period in the list of execution periods?
|
||||
return self.cal.minute_to_session_label(dt, direction="none") in \
|
||||
return self.cal.minute_to_session_label(dt) in \
|
||||
self.execution_periods
|
||||
|
||||
|
||||
@@ -448,7 +448,7 @@ class TradingDayOfMonthRule(six.with_metaclass(ABCMeta, StatelessRule)):
|
||||
|
||||
def should_trigger(self, dt):
|
||||
# is this market minute's period in the list of execution periods?
|
||||
return self.cal.minute_to_session_label(dt, direction="none") in \
|
||||
return self.cal.minute_to_session_label(dt) in \
|
||||
self.execution_periods
|
||||
|
||||
@lazyval
|
||||
|
||||
Reference in New Issue
Block a user