From b69590a2f709c70dd14d817d1a6bee0b1bb0e7b0 Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Thu, 19 Dec 2013 15:01:37 -0500 Subject: [PATCH] ENH: Factor out API methods. Add support for algo scripts. This is a step towards the goal of uniting Quantopian scripts and zipline. To make the syntax of zipline identical to Quantopian we break out the API methods (like order) and turn them into functions. To access the algo object we add a thread local reference to the current algorithm that is accessed in the API functions. TradingAlgorithm now takes either a string or two functions (initialize and handle_data) that it executes. Use api method decorator for methods available in algoscript. Ported appropriate algorithm tests from internal code. --- tests/test_algorithm.py | 254 ++++++++++++++++++++++- tests/test_batchtransform.py | 3 +- zipline/__init__.py | 4 + zipline/algorithm.py | 107 ++++++++-- zipline/api.py | 39 ++++ zipline/examples/buyapple.py | 3 + zipline/examples/quantopian_buy_apple.py | 45 ++++ zipline/test_algorithms.py | 119 ++++++++++- zipline/utils/api_support.py | 42 ++++ 9 files changed, 591 insertions(+), 25 deletions(-) create mode 100644 zipline/api.py create mode 100644 zipline/examples/quantopian_buy_apple.py create mode 100644 zipline/utils/api_support.py diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 1792c97a..b603b1d5 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -16,9 +16,12 @@ from unittest import TestCase from datetime import timedelta import numpy as np +from mock import MagicMock from zipline.utils.test_utils import setup_logger import zipline.utils.factory as factory +import zipline.utils.simfactory as simfactory + from zipline.test_algorithms import (TestRegisterTransformAlgorithm, RecordAlgorithm, TestOrderAlgorithm, @@ -28,13 +31,27 @@ from zipline.test_algorithms import (TestRegisterTransformAlgorithm, TestOrderPercentAlgorithm, TestTargetPercentAlgorithm, TestTargetValueAlgorithm, - EmptyPositionsAlgorithm) + EmptyPositionsAlgorithm, + initialize_noop, + handle_data_noop, + initialize_api, + handle_data_api, + noop_algo, + api_algo, + call_all_order_methods, + record_variables, + record_float_magic + ) + +from zipline.utils.test_utils import drain_zipline, assert_single_position from zipline.sources import (SpecificEquityTrades, DataFrameSource, DataPanelSource) from zipline.transforms import MovingAverage from zipline.finance.trading import SimulationParameters +from zipline.utils.api_support import set_algo_instance +from zipline.algorithm import TradingAlgorithm class TestRecordAlgorithm(TestCase): @@ -235,3 +252,238 @@ class TestPositions(TestCase): for i, expected in enumerate(expected_position_count): self.assertEqual(daily_stats.ix[i]['num_positions'], expected) + + +class TestAlgoScript(TestCase): + def setUp(self): + days = 251 + self.sim_params = factory.create_simulation_parameters(num_days=days) + setup_logger(self) + + trade_history = factory.create_trade_history( + 133, + [10.0] * days, + [100] * days, + timedelta(days=1), + self.sim_params + ) + + self.source = SpecificEquityTrades(event_list=trade_history) + self.df_source, self.df = \ + factory.create_test_df_source(self.sim_params) + + self.zipline_test_config = { + 'sid': 0, + } + + def test_noop(self): + algo = TradingAlgorithm(initialize=initialize_noop, + handle_data=handle_data_noop) + algo.run(self.df) + + def test_noop_string(self): + algo = TradingAlgorithm(script=noop_algo) + algo.run(self.df) + + def test_api_calls(self): + algo = TradingAlgorithm(initialize=initialize_api, + handle_data=handle_data_api) + algo.run(self.df) + + def test_api_calls_string(self): + algo = TradingAlgorithm(script=api_algo) + algo.run(self.df) + + def test_fixed_slippage(self): + # verify order -> transaction -> portfolio position. + # -------------- + test_algo = TradingAlgorithm( + script=""" +from zipline.api import (slippage, + commission, + set_slippage, + set_commission, + order, + record) + +def initialize(context): + model = slippage.FixedSlippage(spread=0.10) + set_slippage(model) + set_commission(commission.PerTrade(100.00)) + context.count = 1 + context.incr = 0 + +def handle_data(context, data): + if context.incr < context.count: + order(0, -1000) + record(price=data[0].price) + + context.incr += 1""", + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + # this matches the value in the algotext initialize + # method, and will be used inside assert_single_position + # to confirm we have as many transactions as orders we + # placed. + self.zipline_test_config['order_count'] = 1 + + # self.zipline_test_config['transforms'] = \ + # test_algo.transform_visitor.transforms.values() + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + + output, _ = assert_single_position(self, zipline) + + # confirm the slippage and commission on a sample + # transaction + recorded_price = output[1]['daily_perf']['recorded_vars']['price'] + transaction = output[1]['daily_perf']['transactions'][0] + self.assertEqual(100.0, transaction['commission']) + expected_spread = 0.05 + expected_commish = 0.10 + expected_price = recorded_price - expected_spread - expected_commish + self.assertEqual(expected_price, transaction['price']) + + def test_volshare_slippage(self): + # verify order -> transaction -> portfolio position. + # -------------- + test_algo = TradingAlgorithm( + script=""" +from zipline.api import * + +def initialize(context): + model = slippage.VolumeShareSlippage( + volume_limit=.3, + price_impact=0.05 + ) + set_slippage(model) + set_commission(commission.PerShare(0.02)) + context.count = 2 + context.incr = 0 + +def handle_data(context, data): + if context.incr < context.count: + # order small lots to be sure the + # order will fill in a single transaction + order(0, 5000) + record(price=data[0].price) + record(volume=data[0].volume) + record(incr=context.incr) + context.incr += 1 + """, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 100 + + # 67 will be used inside assert_single_position + # to confirm we have as many transactions as expected. + # The algo places 2 trades of 5000 shares each. The trade + # events have volume ranging from 100 to 950. The volume cap + # of 0.3 limits the trade volume to a range of 30 - 316 shares. + # The spreadsheet linked below calculates the total position + # size over each bar, and predicts 67 txns will be required + # to fill the two orders. The number of bars and transactions + # differ because some bars result in multiple txns. See + # spreadsheet for details: +# https://www.dropbox.com/s/ulrk2qt0nrtrigb/Volume%20Share%20Worksheet.xlsx + self.zipline_test_config['expected_transactions'] = 67 + + # self.zipline_test_config['transforms'] = \ + # test_algo.transform_visitor.transforms.values() + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + output, _ = assert_single_position(self, zipline) + + # confirm the slippage and commission on a sample + # transaction + per_share_commish = 0.02 + perf = output[1] + transaction = perf['daily_perf']['transactions'][0] + commish = transaction['amount'] * per_share_commish + self.assertEqual(commish, transaction['commission']) + self.assertEqual(2.029, transaction['price']) + + def test_algo_record_vars(self): + test_algo = TradingAlgorithm( + script=record_variables, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + output, _ = drain_zipline(self, zipline) + self.assertEqual(len(output), 252) + incr = [] + for o in output[:200]: + incr.append(o['daily_perf']['recorded_vars']['incr']) + + np.testing.assert_array_equal(incr, range(1, 201)) + + def test_algo_record_allow_mock(self): + """ + Test that values from "MagicMock"ed methods can be passed to record. + + Relevant for our basic/validation and methods like history, which + will end up returning a MagicMock instead of a DataFrame. + """ + test_algo = TradingAlgorithm( + script=record_variables, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + test_algo.record(foo=MagicMock()) + + def _algo_record_float_magic_should_pass(self, var_type): + test_algo = TradingAlgorithm( + script=record_float_magic % var_type, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + output, _ = drain_zipline(self, zipline) + self.assertEqual(len(output), 252) + incr = [] + for o in output[:200]: + incr.append(o['daily_perf']['recorded_vars']['data']) + np.testing.assert_array_equal(incr, [np.nan] * 200) + + def test_algo_record_nan(self): + self._algo_record_float_magic_should_pass('nan') + + def test_order_methods(self): + """Only test that order methods can be called without error. + Correct filling of orders is tested in zipline. + """ + test_algo = TradingAlgorithm( + script=call_all_order_methods, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + + output, _ = drain_zipline(self, zipline) diff --git a/tests/test_batchtransform.py b/tests/test_batchtransform.py index 8f5bdbbe..ece1c2c8 100644 --- a/tests/test_batchtransform.py +++ b/tests/test_batchtransform.py @@ -27,9 +27,10 @@ from zipline.utils.test_utils import setup_logger from zipline.sources.data_source import DataSource import zipline.utils.factory as factory +from zipline.transforms import batch_transform + from zipline.test_algorithms import (BatchTransformAlgorithm, BatchTransformAlgorithmMinute, - batch_transform, ReturnPriceBatchTransform) from zipline.algorithm import TradingAlgorithm diff --git a/zipline/__init__.py b/zipline/__init__.py index 222e4ce3..250d4537 100644 --- a/zipline/__init__.py +++ b/zipline/__init__.py @@ -26,13 +26,17 @@ from . import data from . import finance from . import gens from . import utils +from . import transforms from . algorithm import TradingAlgorithm +from . import api __all__ = [ 'data', 'finance', 'gens', 'utils', + 'transforms', + 'api', 'TradingAlgorithm' ] diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 8d5baf77..01b1fc54 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -22,7 +22,7 @@ from datetime import datetime from itertools import groupby from six.moves import filter -from six import iteritems +from six import iteritems, exec_ from operator import attrgetter from zipline.errors import ( @@ -34,6 +34,7 @@ from zipline.errors import ( from zipline.finance.performance import PerformanceTracker from zipline.sources import DataFrameSource, DataPanelSource from zipline.utils.factory import create_simulation_parameters +from zipline.utils.api_support import set_algo_instance, api_method from zipline.transforms.utils import StatefulTransform from zipline.finance.slippage import ( VolumeShareSlippage, @@ -64,19 +65,21 @@ class TradingAlgorithm(object): A new algorithm could look like this: ``` - class MyAlgo(TradingAlgorithm): - def initialize(self, sids, amount): - self.sids = sids - self.amount = amount + from zipline.api import order - def handle_data(self, data): - sid = self.sids[0] - amount = self.amount - self.order(sid, amount) + def initialize(context): + context.sid = 'AAPL' + context.amount = 100 + + def handle_data(self, data): + sid = context.sid + amount = context.amount + order(sid, amount) ``` - To then to run this algorithm: + To then to run this algorithm pass these functions to + TradingAlgorithm: - my_algo = MyAlgo([0], 100) # first argument has to be list of sids + my_algo = TradingAlgorithm(initialize, handle_data) stats = my_algo.run(data) """ @@ -84,6 +87,16 @@ class TradingAlgorithm(object): """Initialize sids and other state variables. :Arguments: + :Optional: + initialize : function + Function that is called with a single + argument at the begninning of the simulation. + handle_data : function + Function that is called with 2 arguments + (context and data) on every bar. + script : str + Algoscript that contains initialize and + handle_data function definition. data_frequency : str (daily, hourly or minutely) The duration of the bars. annualizer : int @@ -137,13 +150,46 @@ class TradingAlgorithm(object): self.portfolio_needs_update = True self._portfolio = None + # If string is passed in, execute and get reference to + # functions. + self.algoscript = kwargs.pop('script', None) + + if self.algoscript is not None: + self.ns = {} + exec_(self.algoscript, self.ns) + if 'initialize' not in self.ns: + raise ValueError('You must define an initialze function.') + if 'handle_data' not in self.ns: + raise ValueError('You must define a handle_data function.') + self._initialize = self.ns['initialize'] + self._handle_data = self.ns['handle_data'] + + # If two functions are passed in assume initialize and + # handle_data are passed in. + elif kwargs.get('initialize', False) and kwargs.get('handle_data'): + if self.algoscript is not None: + raise ValueError('You can not set script and \ + initialize/handle_data.') + self._initialize = kwargs.pop('initialize') + self._handle_data = kwargs.pop('handle_data') + # an algorithm subclass needs to set initialized to True when # it is fully initialized. self.initialized = False - # call to user-defined constructor method self.initialize(*args, **kwargs) + def initialize(self, *args, **kwargs): + # store algo reference in global space + set_algo_instance(self) + try: + self._initialize(self) + finally: + set_algo_instance(None) + + def handle_data(self, data): + self._handle_data(self, data) + def __repr__(self): """ N.B. this does not yet represent a string that can be used @@ -244,9 +290,6 @@ class TradingAlgorithm(object): """ return self._create_generator(self.sim_params) - def initialize(self, *args, **kwargs): - pass - # TODO: make a new subclass, e.g. BatchAlgorithm, and move # the run method to the subclass, and refactor to put the # generator creation logic into get_generator. @@ -324,14 +367,21 @@ class TradingAlgorithm(object): # create transforms and zipline self.gen = self._create_generator(sim_params) - # loop through simulated_trading, each iteration returns a - # perf dictionary - perfs = [] - for perf in self.gen: - perfs.append(perf) + # store algo reference in global space + set_algo_instance(self) - # convert perf dict to pandas dataframe - daily_stats = self._create_daily_stats(perfs) + try: + # loop through simulated_trading, each iteration returns a + # perf dictionary + perfs = [] + for perf in self.gen: + perfs.append(perf) + + # convert perf dict to pandas dataframe + daily_stats = self._create_daily_stats(perfs) + finally: + # remove algo from global space + set_algo_instance(None) return daily_stats @@ -375,6 +425,7 @@ class TradingAlgorithm(object): 'args': args, 'kwargs': kwargs} + @api_method def record(self, **kwargs): """ Track and record local variable (i.e. attributes) each day. @@ -382,9 +433,11 @@ class TradingAlgorithm(object): for name, value in kwargs.items(): self._recorded_vars[name] = value + @api_method def order(self, sid, amount, limit_price=None, stop_price=None): return self.blotter.order(sid, amount, limit_price, stop_price) + @api_method def order_value(self, sid, value, limit_price=None, stop_price=None): """ Place an order by desired value rather than desired number of shares. @@ -438,6 +491,7 @@ class TradingAlgorithm(object): "Algorithm expects a utc datetime" self.datetime = dt + @api_method def get_datetime(self): """ Returns a copy of the datetime. @@ -454,6 +508,7 @@ class TradingAlgorithm(object): """ self.blotter.transact = transact + @api_method def set_slippage(self, slippage): if not isinstance(slippage, SlippageModel): raise UnsupportedSlippageModel() @@ -461,6 +516,7 @@ class TradingAlgorithm(object): raise OverrideSlippagePostInit() self.slippage = slippage + @api_method def set_commission(self, commission): if not isinstance(commission, (PerShare, PerTrade, PerDollar)): raise UnsupportedCommissionModel() @@ -482,6 +538,7 @@ class TradingAlgorithm(object): self.data_frequency = data_frequency self.annualizer = ANNUALIZER[self.data_frequency] + @api_method def order_percent(self, sid, percent, limit_price=None, stop_price=None): """ Place an order in the specified security corresponding to the given @@ -492,6 +549,7 @@ class TradingAlgorithm(object): value = self.portfolio.portfolio_value * percent return self.order_value(sid, value, limit_price, stop_price) + @api_method def order_target(self, sid, target, limit_price=None, stop_price=None): """ Place an order to adjust a position to a target number of shares. If @@ -507,6 +565,7 @@ class TradingAlgorithm(object): else: return self.order(sid, target, limit_price, stop_price) + @api_method def order_target_value(self, sid, target, limit_price=None, stop_price=None): """ @@ -525,6 +584,7 @@ class TradingAlgorithm(object): else: return self.order_value(sid, target, limit_price, stop_price) + @api_method def order_target_percent(self, sid, target, limit_price=None, stop_price=None): """ @@ -547,6 +607,7 @@ class TradingAlgorithm(object): req_value = target_value - current_value return self.order_value(sid, req_value, limit_price, stop_price) + @api_method def get_open_orders(self, sid=None): if sid is None: return {key: [order.to_api_obj() for order in orders] @@ -557,10 +618,12 @@ class TradingAlgorithm(object): return [order.to_api_obj() for order in orders] return [] + @api_method def get_order(self, order_id): if order_id in self.blotter.orders: return self.blotter.orders[order_id].to_api_obj() + @api_method def cancel_order(self, order_param): order_id = order_param if isinstance(order_param, zipline.protocol.Order): diff --git a/zipline/api.py b/zipline/api.py new file mode 100644 index 00000000..bc012ab0 --- /dev/null +++ b/zipline/api.py @@ -0,0 +1,39 @@ +# +# Copyright 2013 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. + +# Note that part of the API is implemented in TradingAlgorithm as +# methods (e.g. order). These are added to this namespace via the +# decorator `api_methods` inside of algorithm.py. + +import zipline +from .finance import (commission, slippage) +from .utils import math_utils + +from zipline.finance.slippage import ( + FixedSlippage, + VolumeShareSlippage, +) + + +batch_transform = zipline.transforms.BatchTransform + +__all__ = [ + 'slippage', + 'commission', + 'math_utils', + 'batch_transform', + 'FixedSlippage', + 'VolumeShareSlippage' +] diff --git a/zipline/examples/buyapple.py b/zipline/examples/buyapple.py index 5f0a931c..e4ba859d 100755 --- a/zipline/examples/buyapple.py +++ b/zipline/examples/buyapple.py @@ -26,6 +26,9 @@ class BuyApple(TradingAlgorithm): # inherit from TradingAlgorithm """This is the simplest possible algorithm that does nothing but buy 1 apple share on each event. """ + def initialize(self): + pass + def handle_data(self, data): # overload handle_data() method self.order('AAPL', 1) # order SID (=0) and amount (=1 shares) diff --git a/zipline/examples/quantopian_buy_apple.py b/zipline/examples/quantopian_buy_apple.py new file mode 100644 index 00000000..bf0d69a2 --- /dev/null +++ b/zipline/examples/quantopian_buy_apple.py @@ -0,0 +1,45 @@ +# +# Copyright 2013 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 datetime import datetime +import pytz + +from zipline import TradingAlgorithm +from zipline.utils.factory import load_from_yahoo + +from zipline.api import order + + +def initialize(context): + context.test = 10 + + +def handle_date(context, data): + order('AAPL', 10) + print(context.test) + + +if __name__ == '__main__': + import pylab as pl + start = datetime(2008, 1, 1, 0, 0, 0, 0, pytz.utc) + end = datetime(2010, 1, 1, 0, 0, 0, 0, pytz.utc) + data = load_from_yahoo(stocks=['AAPL'], indexes={}, start=start, + end=end) + data = data.dropna() + algo = TradingAlgorithm(initialize=initialize, + handle_data=handle_date) + results = algo.run(data) + results.portfolio_value.plot() + pl.show() diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index f9a958e2..614476ea 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -78,7 +78,7 @@ from six.moves import range from six import itervalues from zipline.algorithm import TradingAlgorithm -from zipline.finance.slippage import FixedSlippage +from zipline.api import FixedSlippage class TestAlgorithm(TradingAlgorithm): @@ -665,3 +665,120 @@ class EmptyPositionsAlgorithm(TradingAlgorithm): # Should be 0 when all positions are exited. self.record(num_positions=len(self.portfolio.positions)) + + +############################## +# Quantopian style algorithms +from zipline.api import (order, + set_slippage, + record) + + +# Noop algo +def initialize_noop(context): + pass + + +def handle_data_noop(context, data): + pass + + +# API functions +def initialize_api(context): + context.incr = 0 + context.sale_price = None + set_slippage(FixedSlippage()) + + +def handle_data_api(context, data): + if context.incr == 0: + assert 0 not in context.portfolio.positions + else: + assert context.portfolio.positions[0]['amount'] == \ + context.incr, "Orders not filled immediately." + assert context.portfolio.positions[0]['last_sale_price'] == \ + data[0].price, "Orders not filled at current price." + context.incr += 1 + order(0, 1) + + record(incr=context.incr) + +########################### +# AlgoScripts as strings +noop_algo = """ +# Noop algo +def initialize(context): + pass + +def handle_data(context, data): + pass +""" + +api_algo = """ +from zipline.api import (order, + set_slippage, + FixedSlippage, + record) + +def initialize(context): + context.incr = 0 + context.sale_price = None + set_slippage(FixedSlippage()) + +def handle_data(context, data): + if context.incr == 0: + assert 0 not in context.portfolio.positions + else: + assert context.portfolio.positions[0]['amount'] == \ + context.incr, "Orders not filled immediately." + assert context.portfolio.positions[0]['last_sale_price'] == \ + data[0].price, "Orders not filled at current price." + context.incr += 1 + order(0, 1) + + record(incr=context.incr) +""" + +call_all_order_methods = """ +from zipline.api import (order, + order_value, + order_percent, + order_target, + order_target_value, + order_target_percent) + +def initialize(context): + pass + +def handle_data(context, data): + order(0, 10) + order_value(0, 300) + order_percent(0, .1) + order_target(0, 100) + order_target_value(0, 100) + order_target_percent(0, .2) +""" + +record_variables = """ +from zipline.api import record + +def initialize(context): + context.stocks = [0, 1] + context.incr = 0 + +def handle_data(context, data): + context.incr += 1 + record(incr=context.incr) +""" + +record_float_magic = """ +from zipline.api import record + +def initialize(context): + context.stocks = [0, 1] + context.incr = 0 + +def handle_data(context, data): + context.incr += 1 + record(data=float('%s')) +""" diff --git a/zipline/utils/api_support.py b/zipline/utils/api_support.py new file mode 100644 index 00000000..0a19ca1c --- /dev/null +++ b/zipline/utils/api_support.py @@ -0,0 +1,42 @@ +# +# Copyright 2013 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 functools import wraps +import zipline.api + +import threading +context = threading.local() + + +def get_algo_instance(): + return getattr(context, 'algorithm', None) + + +def set_algo_instance(algo): + context.algorithm = algo + + +def api_method(f): + # Decorator that adds the decorated class method as a callable + # function (wrapped) to zipline.api + @wraps(f) + def wrapped(*args, **kwargs): + # Get the instance and call the method + return getattr(get_algo_instance(), f.__name__)(*args, **kwargs) + # Add functor to zipline.api + setattr(zipline.api, f.__name__, wrapped) + zipline.api.__all__.append(f.__name__) + + return f