ENH: Adds a new api_method called get_environment so that users may

check if their algorithm is running in zipline or on Quantopian.
This commit is contained in:
llllllllll
2014-08-07 16:23:21 -04:00
parent 4420e370ff
commit 7563124547
3 changed files with 26 additions and 0 deletions
+8
View File
@@ -53,6 +53,7 @@ from zipline.test_algorithms import (
SetMaxOrderCountAlgorithm,
SetMaxOrderSizeAlgorithm,
api_algo,
api_get_environment_algo,
api_symbol_algo,
call_all_order_methods,
call_order_in_init,
@@ -407,6 +408,13 @@ class TestAlgoScript(TestCase):
algo = TradingAlgorithm(script=api_algo)
algo.run(self.df)
def test_api_get_environment(self):
environment = 'zipline'
algo = TradingAlgorithm(script=api_get_environment_algo,
environment=environment)
algo.run(self.df)
self.assertEqual(algo.environment, environment)
def test_api_symbol(self):
algo = TradingAlgorithm(script=api_symbol_algo)
algo.run(self.df)
+8
View File
@@ -126,6 +126,8 @@ class TradingAlgorithm(object):
How much capital to start with.
instant_fill : bool <default: False>
Whether to fill orders immediately or on next bar.
environment : str <default: 'zipline'>
The environment that this algorithm is running in.
"""
self.datetime = None
@@ -139,6 +141,8 @@ class TradingAlgorithm(object):
self._recorded_vars = {}
self.namespace = kwargs.get('namespace', {})
self._environment = kwargs.pop('environment', 'zipline')
self.logger = None
self.benchmark_return_source = None
@@ -470,6 +474,10 @@ class TradingAlgorithm(object):
'args': args,
'kwargs': kwargs}
@api_method
def get_environment(self):
return self._environment
@api_method
def record(self, *args, **kwargs):
"""
+10
View File
@@ -939,6 +939,16 @@ def handle_data(context, data):
record(incr=context.incr)
"""
api_get_environment_algo = """
from zipline.api import get_environment, order, symbol
def initialize(context):
context.environment = get_environment()
handle_data = lambda context, data: order(symbol(0), 1)
"""
api_symbol_algo = """
from zipline.api import (order,
symbol)