added CLOSE_POSITION as source type, added pt.close_position_event(), added process_close_postion(), added close processing to tradesimulation, added unittest for close_position_event

tests/test_perf_tracking.py

@@ -46,7 +46,8 @@ from zipline.finance.commission import PerShare, PerTrade, PerDollar
 from zipline.finance import trading
 from zipline.utils.factory import create_random_simulation_parameters
 import zipline.protocol as zp
-from zipline.protocol import Event
+from zipline.protocol import Event, DATASOURCE_TYPE
+from zipline.sources.data_frame_source import DataPanelSource
 
 logger = logging.getLogger('Test Perf Tracking')
 
@@ -1937,6 +1938,37 @@ class TestPerformanceTracker(unittest.TestCase):
 
             check_perf_tracker_serialization(tracker)
 
+    def test_close_position_event(self):
+        pt = perf.PositionTracker()
+        dt = pd.Timestamp("1984/03/06 3:00PM")
+        pos1 = perf.Position(1, amount=np.float64(120.0),
+                             last_sale_date=dt, last_sale_price=3.4)
+        pos2 = perf.Position(2, amount=np.float64(-100.0),
+                             last_sale_date=dt, last_sale_price=3.4)
+        pt.update_positions({1: pos1, 2: pos2})
+
+        event_type = DATASOURCE_TYPE.CLOSE_POSITION
+        index = [dt + timedelta(days=1)]
+        pan = pd.Panel({1: pd.DataFrame({'price': 1, 'volume': 0,
+                                         'type': event_type}, index=index),
+                        2: pd.DataFrame({'price': 1, 'volume': 0,
+                                         'type': event_type}, index=index),
+                        3: pd.DataFrame({'price': 1, 'volume': 0,
+                                         'type': event_type}, index=index)})
+
+        source = DataPanelSource(pan)
+        for i, event in enumerate(source):
+            txn = pt.create_close_position_transaction(event)
+            if event.sid == 1:
+                # Test owned long
+                self.assertEqual(-120, txn.amount)
+            elif event.sid == 2:
+                # Test owned short
+                self.assertEqual(100, txn.amount)
+            elif event.sid == 3:
+                # Test not-owned SID
+                self.assertIsNone(txn)
+
     def test_serialization(self):
         start_dt = datetime(year=2008,
                             month=10,

zipline/finance/performance/position_tracker.py

@@ -13,6 +13,7 @@ except ImportError:
 from six import iteritems
 from six.moves import map, filter
 
+from zipline.finance.slippage import Transaction
 from zipline.utils.serialization_utils import (
     VERSION_LABEL
 )
@@ -217,6 +218,19 @@ class PositionTracker(object):
         net_cash_payment = payments['cash_amount'].fillna(0).sum()
         return net_cash_payment
 
+    def create_close_position_transaction(self, event):
+        if not self._position_amounts.get(event.sid):
+            return None
+        txn = Transaction(
+            sid=event.sid,
+            amount=(-1 * self._position_amounts[event.sid]),
+            dt=event.dt,
+            price=event.price,
+            commission=0,
+            order_id=0
+        )
+        return txn
+
     def get_positions(self):
 
         positions = self._positions_store

zipline/finance/performance/tracker.py

@@ -337,6 +337,10 @@ class PerformanceTracker(object):
 
         self.all_benchmark_returns[midnight] = event.returns
 
+    def process_close_position(self, event):
+        txn = self.position_tracker.create_close_position_transaction(event)
+        self.process_transaction(txn)
+
     def check_upcoming_dividends(self, midnight_of_date_that_just_ended):
         """
         Check if we currently own any stocks with dividends whose ex_date is

zipline/gens/tradesimulation.py

@@ -204,6 +204,8 @@ class AlgorithmSimulator(object):
         perf_process_split = self.algo.perf_tracker.process_split
         perf_process_dividend = self.algo.perf_tracker.process_dividend
         perf_process_commission = self.algo.perf_tracker.process_commission
+        perf_process_close_position = \
+            self.algo.perf_tracker.process_close_position
         blotter_process_trade = self.algo.blotter.process_trade
         blotter_process_benchmark = self.algo.blotter.process_benchmark
 
@@ -219,6 +221,7 @@ class AlgorithmSimulator(object):
         # custom events.
         trades = []
         customs = []
+        closes = []
 
         # splits and dividends are processed once a day.
         #
@@ -247,6 +250,8 @@ class AlgorithmSimulator(object):
                 if dividends is None:
                     dividends = []
                 dividends.append(event)
+            elif event.type == DATASOURCE_TYPE.CLOSE_POSITION:
+                closes.append(event)
             else:
                 raise log.warn("Unrecognized event=%s".format(event))
 
@@ -282,6 +287,10 @@ class AlgorithmSimulator(object):
         for custom in customs:
             self.update_universe(custom)
 
+        for close in closes:
+            self.update_universe(close)
+            perf_process_close_position(close)
+
         if splits is not None:
             for split in splits:
                 # process_split is not assigned to a variable since it is

zipline/protocol.py

@@ -42,7 +42,8 @@ DATASOURCE_TYPE = Enum(
     'DONE',
     'CUSTOM',
     'BENCHMARK',
-    'COMMISSION'
+    'COMMISSION',
+    'CLOSE_POSITION'
 )
 
 # Expected fields/index values for a dividend Series.