diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index d01adca1..9bd7dc5a 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -356,3 +356,17 @@ class DailyReturn(object): def __repr__(self): return str(self.date) + " - " + str(self.returns) + + +class use_environment(object): + """A decorator to wrap a method in a particular + trading environment.""" + + def __init__(self, environment): + self.env = environment + + def __call__(self, func): + def wrapper(*args, **kwargs): + with self.env: + return func(*args, **kwargs) + return wrapper