MAINT: Support the passing of a time rule positionally on the date_rule arg

But log a warning to the user
This commit is contained in:
Andrew Liang
2016-05-20 14:46:47 -04:00
parent cf1687ec72
commit 1056501b27
3 changed files with 82 additions and 2 deletions
+67
View File
@@ -16,6 +16,7 @@ from collections import namedtuple
import datetime
from datetime import timedelta
from textwrap import dedent
import warnings
from unittest import TestCase, skip
from copy import deepcopy
@@ -1854,6 +1855,72 @@ def handle_data(context, data):
)
algo.run(self.data_portal)
def test_schedule_function_time_rule_positionally_misplaced(self):
"""
Test that when a user specifies a time rule for the date_rule argument,
but no rule in the time_rule argument
(e.g. schedule_function(func, <time_rule>)), we assume that means
assign a time rule but no date rule
"""
sim_params = factory.create_simulation_parameters(
start=pd.Timestamp('2006-01-12', tz='UTC'),
end=pd.Timestamp('2006-01-13', tz='UTC'),
data_frequency='minute'
)
algocode = dedent("""
from zipline.api import time_rules, schedule_function
def do_at_open(context, data):
context.done_at_open.append(context.get_datetime())
def do_at_close(context, data):
context.done_at_close.append(context.get_datetime())
def initialize(context):
context.done_at_open = []
context.done_at_close = []
schedule_function(do_at_open, time_rules.market_open())
schedule_function(do_at_close, time_rules.market_close())
def handle_data(algo, data):
pass
""")
with warnings.catch_warnings(record=True) as w:
algo = TradingAlgorithm(
script=algocode,
sim_params=sim_params,
env=self.env
)
algo.run(self.data_portal)
self.assertEqual(len(w), 2)
for i, warning in enumerate(w):
self.assertIsInstance(warning.message, UserWarning)
self.assertEqual(
warning.message.args[0],
'Got a time rule for the second positional argument '
'date_rule. You should use keyword argument '
'time_rule= when calling schedule_function without '
'specifying a date_rule'
)
# The warnings come from line 13 and 14 in the algocode
self.assertEqual(warning.lineno, 13 + i)
self.assertEqual(
algo.done_at_open,
[pd.Timestamp('2006-01-12 14:31:00', tz='UTC'),
pd.Timestamp('2006-01-13 14:31:00', tz='UTC')]
)
self.assertEqual(
algo.done_at_close,
[pd.Timestamp('2006-01-12 20:59:00', tz='UTC'),
pd.Timestamp('2006-01-13 20:59:00', tz='UTC')]
)
class TestCapitalChanges(WithLogger,
WithDataPortal,
+14 -2
View File
@@ -100,6 +100,8 @@ from zipline.utils.events import (
make_eventrule,
date_rules,
time_rules,
AfterOpen,
BeforeClose
)
from zipline.utils.factory import create_simulation_parameters
from zipline.utils.math_utils import (
@@ -951,11 +953,21 @@ class TradingAlgorithm(object):
:class:`zipline.api.date_rules`
:class:`zipline.api.time_rules`
"""
# When the user calls schedule_function(func, <time_rule>), assume that
# the user meant to specify a time rule but no date rule, instead of
# a date rule and no time rule as the signature suggests
if isinstance(date_rule, (AfterOpen, BeforeClose)) and not time_rule:
warnings.warn('Got a time rule for the second positional argument '
'date_rule. You should use keyword argument '
'time_rule= when calling schedule_function without '
'specifying a date_rule', stacklevel=3)
date_rule = date_rule or date_rules.every_day()
time_rule = ((time_rule or time_rules.market_open())
time_rule = ((time_rule or time_rules.every_minute())
if self.sim_params.data_frequency == 'minute' else
# If we are in daily mode the time_rule is ignored.
zipline.utils.events.Always())
time_rules.every_minute())
self.add_event(
make_eventrule(date_rule, time_rule, half_days),
+1
View File
@@ -689,6 +689,7 @@ class date_rules(object):
class time_rules(object):
market_open = AfterOpen
market_close = BeforeClose
every_minute = Always
def make_eventrule(date_rule, time_rule, half_days=True):