ENH: Adds a new api method schedule_function.

schedule_function takes a date rule, a time rule, and a function and
will call the function, passing context and data only when the two rules
fire. This allows for code that is conditional to the datetime of the
algo.

This is implemented internally with `Event` objects which are pairings
of `EventRule`s and callbacks.

handle_data becomes a special event with a rule that always fires. This
makes the logic for handling events more complete and compact.
This commit is contained in:
Joe Jevnik
2014-10-06 13:42:36 -04:00
parent 6050b57fa4
commit 3c37704a5b
10 changed files with 1585 additions and 7 deletions
+40 -2
View File
@@ -12,7 +12,7 @@
# 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.
import datetime
from datetime import timedelta
from mock import MagicMock
from six.moves import range
@@ -77,6 +77,7 @@ from zipline.transforms import MovingAverage
from zipline.finance.execution import LimitOrder
from zipline.finance.trading import SimulationParameters
from zipline.utils.api_support import set_algo_instance
from zipline.utils.events import DateRuleFactory, TimeRuleFactory
from zipline.algorithm import TradingAlgorithm
@@ -174,6 +175,44 @@ class TestMiscellaneousAPI(TestCase):
sim_params=self.sim_params)
algo.run(self.source)
def test_schedule_function(self):
date_rules = DateRuleFactory
time_rules = TimeRuleFactory
def incrementer(algo, data):
algo.func_called += 1
self.assertEqual(
algo.get_datetime().time(),
datetime.time(hour=14, minute=31),
)
def initialize(algo):
algo.func_called = 0
algo.days = 1
algo.date = None
algo.schedule_function(
func=incrementer,
date_rule=date_rules.every_day(),
time_rule=time_rules.market_open(),
)
def handle_data(algo, data):
if not algo.date:
algo.date = algo.get_datetime().date()
if algo.date < algo.get_datetime().date():
algo.days += 1
algo.date = algo.get_datetime().date()
algo = TradingAlgorithm(
initialize=initialize,
handle_data=handle_data,
sim_params=self.sim_params,
)
algo.run(self.source)
self.assertEqual(algo.func_called, algo.days)
class TestTransformAlgorithm(TestCase):
def setUp(self):
@@ -840,7 +879,6 @@ class TestTradingControls(TestCase):
self.check_algo_succeeds(algo, handle_data, order_count=20)
def test_long_only(self):
# Sell immediately -> fail immediately.
def handle_data(algo, data):
algo.order(self.sid, -1)