From 64eed84bffcafc0b92ac67db021b1f38b11e0cc7 Mon Sep 17 00:00:00 2001
From: Delaney Granizo-Mackenzie <Captain.Kanuk@gmail.com>
Date: Wed, 25 Feb 2015 18:45:46 -0500
Subject: [PATCH] MAINT: Added pickle protocol methods into zipline.

Added pickle support to many zipline methods. This will enable
them to be serialized.
---
 zipline/finance/blotter.py              | 25 ++++++++++++-
 zipline/finance/commission.py           |  8 ++--
 zipline/finance/performance/period.py   | 29 +++++++++++++-
 zipline/finance/performance/position.py |  7 +++-
 zipline/finance/performance/tracker.py  | 26 ++++++++++++-
 zipline/finance/risk/cumulative.py      | 18 ++++++++-
 zipline/finance/risk/period.py          | 16 +++++++-
 zipline/finance/risk/report.py          | 12 +++++-
 zipline/finance/slippage.py             | 22 +++++++++--
 zipline/protocol.py                     | 20 ++++++----
 zipline/utils/serialization_utils.py    | 50 +++++++++++++++++++++++++
 11 files changed, 211 insertions(+), 22 deletions(-)
 create mode 100644 zipline/utils/serialization_utils.py

diff --git a/zipline/finance/blotter.py b/zipline/finance/blotter.py
index 00501bcd..823b3358 100644
--- a/zipline/finance/blotter.py
+++ b/zipline/finance/blotter.py
@@ -35,6 +35,7 @@ from zipline.finance.commission import PerShare
 log = Logger('Blotter')
 
 from zipline.utils.protocol_utils import Enum
+from zipline.utils.serialization_utils import SerializeableZiplineObject
 
 ORDER_STATUS = Enum(
     'OPEN',
@@ -45,7 +46,7 @@ ORDER_STATUS = Enum(
 )
 
 
-class Blotter(object):
+class Blotter(SerializeableZiplineObject):
 
     def __init__(self):
         self.transact = transact_partial(VolumeShareSlippage(), PerShare())
@@ -247,8 +248,21 @@ class Blotter(object):
 
             yield txn, order
 
+    def __getstate__(self):
 
-class Order(object):
+        state_to_save = ['new_orders', 'orders', '_status']
+
+        state_dict = {k: self.__dict__[k] for k in state_to_save
+                      if k in self.__dict__}
+
+        # Have to handle defaultdicts specially
+        state_dict['open_orders'] = \
+            self._defaultdict_list_get_state(self.open_orders)
+
+        return state_dict
+
+
+class Order(SerializeableZiplineObject):
     def __init__(self, dt, sid, amount, stop=None, limit=None, filled=0,
                  commission=None, id=None):
         """
@@ -385,3 +399,10 @@ class Order(object):
         Unicode representation for this object.
         """
         return text_type(repr(self))
+
+    def __getstate__(self):
+        state_dict = super(Order, self).__getstate__()
+
+        state_dict['_status'] = self._status
+
+        return state_dict
diff --git a/zipline/finance/commission.py b/zipline/finance/commission.py
index 0c064722..f7bd84c7 100644
--- a/zipline/finance/commission.py
+++ b/zipline/finance/commission.py
@@ -13,8 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from zipline.utils.serialization_utils import SerializeableZiplineObject
 
-class PerShare(object):
+
+class PerShare(SerializeableZiplineObject):
     """
     Calculates a commission for a transaction based on a per
     share cost with an optional minimum cost per trade.
@@ -51,7 +53,7 @@ class PerShare(object):
             return abs(commission / transaction.amount), commission
 
 
-class PerTrade(object):
+class PerTrade(SerializeableZiplineObject):
     """
     Calculates a commission for a transaction based on a per
     trade cost.
@@ -78,7 +80,7 @@ class PerTrade(object):
         return abs(self.cost / transaction.amount), self.cost
 
 
-class PerDollar(object):
+class PerDollar(SerializeableZiplineObject):
     """
     Calculates a commission for a transaction based on a per
     dollar cost.
diff --git a/zipline/finance/performance/period.py b/zipline/finance/performance/period.py
index 0951989d..6a9ccaee 100644
--- a/zipline/finance/performance/period.py
+++ b/zipline/finance/performance/period.py
@@ -92,11 +92,13 @@ from six import iteritems, itervalues
 import zipline.protocol as zp
 from . position import positiondict
 
+from zipline.utils.serialization_utils import SerializeableZiplineObject
+
 log = logbook.Logger('Performance')
 TRADE_TYPE = zp.DATASOURCE_TYPE.TRADE
 
 
-class PerformancePeriod(object):
+class PerformancePeriod(SerializeableZiplineObject):
 
     def __init__(
             self,
@@ -574,3 +576,28 @@ class PerformancePeriod(object):
             if pos.amount != 0:
                 positions.append(pos.to_dict())
         return positions
+
+    def __getstate__(self):
+        state_dict = super(PerformancePeriod, self).__getstate__()
+
+        state_dict['_portfolio_store'] = self._portfolio_store
+        state_dict['_account_store'] = self._account_store
+
+        # We need to handle the defaultdict specially, otherwise
+        # msgpack will unpack it as a dict, causing KeyError
+        # nastiness.
+        state_dict['processed_transactions'] = \
+            self._defaultdict_list_get_state(self.processed_transactions)
+        state_dict['orders_by_modified'] = \
+            self._defaultdict_ordered_get_state(self.orders_by_modified)
+        state_dict['positions'] = \
+            self._positiondict_get_state(self.positions)
+        state_dict['_positions_store'] = \
+            self._positions_get_state(self._positions_store)
+
+        return state_dict
+
+    def __setstate__(self, state):
+        super(PerformancePeriod, self).__setstate__(state)
+
+        self.initialize_position_calc_arrays()
diff --git a/zipline/finance/performance/position.py b/zipline/finance/performance/position.py
index 52bda2be..b7f2a2aa 100644
--- a/zipline/finance/performance/position.py
+++ b/zipline/finance/performance/position.py
@@ -41,10 +41,12 @@ from math import (
 import logbook
 import zipline.protocol as zp
 
+from zipline.utils.serialization_utils import SerializeableZiplineObject
+
 log = logbook.Logger('Performance')
 
 
-class Position(object):
+class Position(SerializeableZiplineObject):
 
     def __init__(self, sid, amount=0, cost_basis=0.0,
                  last_sale_price=0.0, last_sale_date=None):
@@ -207,6 +209,9 @@ last_sale_price: {last_sale_price}"
             'last_sale_price': self.last_sale_price
         }
 
+    def __getstate__(self):
+        return self.__dict__
+
 
 class positiondict(dict):
 
diff --git a/zipline/finance/performance/tracker.py b/zipline/finance/performance/tracker.py
index b1f5a12e..d27ecf2e 100644
--- a/zipline/finance/performance/tracker.py
+++ b/zipline/finance/performance/tracker.py
@@ -59,6 +59,7 @@ Performance Tracking
 
 from __future__ import division
 import logbook
+import pickle
 
 import numpy as np
 import pandas as pd
@@ -70,11 +71,12 @@ from zipline.finance import trading
 from . period import PerformancePeriod
 
 from zipline.finance.trading import with_environment
+from zipline.utils.serialization_utils import SerializeableZiplineObject
 
 log = logbook.Logger('Performance')
 
 
-class PerformanceTracker(object):
+class PerformanceTracker(SerializeableZiplineObject):
     """
     Tracks the performance of the algorithm.
     """
@@ -483,3 +485,25 @@ class PerformanceTracker(object):
 
         risk_dict = self.risk_report.to_dict()
         return risk_dict
+
+    def __getstate__(self):
+        state_dict = super(PerformanceTracker, self).__getstate__()
+
+        state_dict['dividend_frame'] = pickle.dumps(self.dividend_frame)
+
+        state_dict['_dividend_count'] = self._dividend_count
+
+        return state_dict
+
+    def __setstate__(self, state):
+        super(PerformanceTracker, self).__setstate__(state)
+
+        # Handle the dividend frame specially
+        self.dividend_frame = pickle.loads(state['dividend_frame'])
+
+        # We have to restore the references to the objects,
+        # as the perf periods have been reconstructed as different objects
+        # with the same values.
+        self.perf_periods[0] = self.minute_performance
+        self.perf_periods[1] = self.cumulative_performance
+        self.perf_periods[2] = self.todays_performance
diff --git a/zipline/finance/risk/cumulative.py b/zipline/finance/risk/cumulative.py
index 21269ef9..0f5b5129 100644
--- a/zipline/finance/risk/cumulative.py
+++ b/zipline/finance/risk/cumulative.py
@@ -35,6 +35,8 @@ from . risk import (
     sortino_ratio,
 )
 
+from zipline.utils.serialization_utils import SerializeableZiplineObject
+
 log = logbook.Logger('Risk Cumulative')
 
 
@@ -71,7 +73,7 @@ def information_ratio(algo_volatility, algorithm_return, benchmark_return):
         / algo_volatility)
 
 
-class RiskMetricsCumulative(object):
+class RiskMetricsCumulative(SerializeableZiplineObject):
     """
     :Usage:
         Instantiate RiskMetricsCumulative once.
@@ -454,3 +456,17 @@ algorithm_returns ({algo_count}) in range {start} : {end} on {dt}"
         beta = algorithm_covariance / benchmark_variance
 
         return beta
+
+    def __getstate__(self):
+        state_dict = \
+            {k: v for k, v in self.__dict__.iteritems() if
+                (not k.startswith('_') and not k == 'treasury_curves')}
+
+        return state_dict
+
+    def __setstate__(self, state):
+        super(RiskMetricsCumulative, self).__setstate__(state)
+
+        # This are big and we don't need to serialize them
+        # pop them back in now
+        self.treasury_curves = trading.environment.treasury_curves
diff --git a/zipline/finance/risk/period.py b/zipline/finance/risk/period.py
index c686cdae..51762e85 100644
--- a/zipline/finance/risk/period.py
+++ b/zipline/finance/risk/period.py
@@ -36,13 +36,15 @@ from . risk import (
     sortino_ratio,
 )
 
+from zipline.utils.serialization_utils import SerializeableZiplineObject
+
 log = logbook.Logger('Risk Period')
 
 choose_treasury = functools.partial(risk.choose_treasury,
                                     risk.select_treasury_duration)
 
 
-class RiskMetricsPeriod(object):
+class RiskMetricsPeriod(SerializeableZiplineObject):
     def __init__(self, start_date, end_date, returns,
                  benchmark_returns=None):
 
@@ -304,3 +306,15 @@ class RiskMetricsPeriod(object):
             return 0.0
 
         return 1.0 - math.exp(max_drawdown)
+
+    def __getstate__(self):
+        state_dict = \
+            {k: v for k, v in self.__dict__.iteritems() if
+             (not k.startswith('_') and not k == 'treasury_curves')}
+
+        return state_dict
+
+    def __setstate__(self, state):
+        super(RiskMetricsPeriod, self).__setstate__(state)
+
+        self.treasury_curves = trading.environment.treasury_curves
diff --git a/zipline/finance/risk/report.py b/zipline/finance/risk/report.py
index b965a1a1..c7f73657 100644
--- a/zipline/finance/risk/report.py
+++ b/zipline/finance/risk/report.py
@@ -61,10 +61,12 @@ from dateutil.relativedelta import relativedelta
 
 from . period import RiskMetricsPeriod
 
+from zipline.utils.serialization_utils import SerializeableZiplineObject
+
 log = logbook.Logger('Risk Report')
 
 
-class RiskReport(object):
+class RiskReport(SerializeableZiplineObject):
     def __init__(self, algorithm_returns, sim_params, benchmark_returns=None):
         """
         algorithm_returns needs to be a list of daily_return objects
@@ -138,3 +140,11 @@ class RiskReport(object):
             cur_start = cur_start + relativedelta(months=1)
 
         return ends
+
+    def __getstate__(self):
+        state_dict = super(RiskReport, self).__getstate__()
+
+        if '_dividend_count' in dir(self):
+            state_dict['_dividend_count'] = self._dividend_count
+
+        return state_dict
diff --git a/zipline/finance/slippage.py b/zipline/finance/slippage.py
index cedcab21..656e9aad 100644
--- a/zipline/finance/slippage.py
+++ b/zipline/finance/slippage.py
@@ -24,6 +24,7 @@ from functools import partial
 from six import with_metaclass
 
 from zipline.protocol import DATASOURCE_TYPE
+from zipline.utils.serialization_utils import SerializeableZiplineObject
 
 SELL = 1 << 0
 BUY = 1 << 1
@@ -109,7 +110,7 @@ def transact_partial(slippage, commission):
     return partial(transact_stub, slippage, commission)
 
 
-class Transaction(object):
+class Transaction(SerializeableZiplineObject):
 
     def __init__(self, sid, amount, dt, price, order_id, commission=None):
         self.sid = sid
@@ -128,6 +129,9 @@ class Transaction(object):
         del py['type']
         return py
 
+    def __getstate__(self):
+        return self.__dict__
+
 
 def create_transaction(event, order, price, amount):
 
@@ -190,7 +194,7 @@ class SlippageModel(with_metaclass(abc.ABCMeta)):
         return self.simulate(event, current_orders, **kwargs)
 
 
-class VolumeShareSlippage(SlippageModel):
+class VolumeShareSlippage(SlippageModel, SerializeableZiplineObject):
 
     def __init__(self,
                  volume_limit=.25,
@@ -246,8 +250,14 @@ class VolumeShareSlippage(SlippageModel):
             math.copysign(cur_volume, order.direction)
         )
 
+    def __getstate__(self):
+        return self.__dict__
 
-class FixedSlippage(SlippageModel):
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+
+class FixedSlippage(SlippageModel, SerializeableZiplineObject):
 
     def __init__(self, spread=0.0):
         """
@@ -264,3 +274,9 @@ class FixedSlippage(SlippageModel):
             event.price + (self.spread / 2.0 * order.direction),
             order.amount,
         )
+
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
diff --git a/zipline/protocol.py b/zipline/protocol.py
index 0fd5c43e..038ab693 100644
--- a/zipline/protocol.py
+++ b/zipline/protocol.py
@@ -22,6 +22,7 @@ from . utils.math_utils import nanstd, nanmean, nansum
 
 from zipline.finance.trading import with_environment
 from zipline.utils.algo_instance import get_algo_instance
+from zipline.utils.serialization_utils import SerializeableZiplineObject
 
 # Datasource type should completely determine the other fields of a
 # message with its type.
@@ -119,7 +120,7 @@ class Order(Event):
     pass
 
 
-class Portfolio(object):
+class Portfolio(SerializeableZiplineObject):
 
     def __init__(self):
         self.capital_used = 0.0
@@ -138,8 +139,11 @@ class Portfolio(object):
     def __repr__(self):
         return "Portfolio({0})".format(self.__dict__)
 
+    def __getstate__(self):
+        return self.__dict__
 
-class Account(object):
+
+class Account(SerializeableZiplineObject):
     '''
     The account object tracks information about the trading account. The
     values are updated as the algorithm runs and its keys remain unchanged.
@@ -171,14 +175,11 @@ class Account(object):
     def __repr__(self):
         return "Account({0})".format(self.__dict__)
 
-    def _get_state(self):
-        return 'Account', self.__dict__
-
-    def _set_state(self, saved_state):
-        self.__dict__.update(saved_state)
+    def __getstate__(self):
+        return self.__dict__
 
 
-class Position(object):
+class Position(SerializeableZiplineObject):
 
     def __init__(self, sid):
         self.sid = sid
@@ -192,6 +193,9 @@ class Position(object):
     def __repr__(self):
         return "Position({0})".format(self.__dict__)
 
+    def __getstate__(self):
+        return self.__dict__
+
 
 class Positions(dict):
 
diff --git a/zipline/utils/serialization_utils.py b/zipline/utils/serialization_utils.py
new file mode 100644
index 00000000..9da19aa8
--- /dev/null
+++ b/zipline/utils/serialization_utils.py
@@ -0,0 +1,50 @@
+class SerializeableZiplineObject(object):
+    """
+    This class implements the basic set and get state methods used for
+    serialization. It also serves as a demarkation of which objects we
+    serialize.
+    """
+
+    def __getstate__(self):
+        """
+        Many get_state methods need this one line of code.
+        This method deduplicates the code calls.
+        """
+        state_dict = \
+            {k: v for k, v in self.__dict__.iteritems()
+                if not k.startswith('_')}
+        return state_dict
+
+    def __setstate__(self, state):
+        """
+        Many objects require only this code.
+        """
+        self.__dict__.update(state)
+
+    # =====================================================
+    # These are helper methods for some problem data types.
+    # =====================================================
+
+    def _defaultdict_list_get_state(self, d):
+        return {
+            '__original.type__': 'encoded.defaultdict_list',
+            'as_dict': dict(d)
+        }
+
+    def _defaultdict_ordered_get_state(self, d):
+        return {
+            '__original.type__': 'encoded.defaultdict_ordered',
+            'as_dict': dict(d)
+        }
+
+    def _positiondict_get_state(self, d):
+        return {
+            '__original.type__': 'encoded.positiondict',
+            'as_dict': dict(d)
+        }
+
+    def _positions_get_state(self, d):
+        return {
+            '__original.type__': 'encoded.Positions',
+            'as_dict': dict(d)
+        }