mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-10 19:26:14 +08:00
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:
@@ -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
@@ -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),
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user