From efa6d8dbce86777e69973f7e078021ace09e9781 Mon Sep 17 00:00:00 2001
From: jfkirk <james.f.kirk@gmail.com>
Date: Thu, 9 Jul 2015 13:28:40 -0400
Subject: [PATCH] ENH: Adds a perf tracker method to handle SIDs leaving the
 universe

---
 tests/test_perf_tracking.py            | 40 ++++++++++++++++++++++++++
 zipline/finance/performance/tracker.py | 16 +++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py
index 088d88dc..676cd14b 100644
--- a/tests/test_perf_tracking.py
+++ b/tests/test_perf_tracking.py
@@ -2020,6 +2020,46 @@ class TestPerformanceTracker(unittest.TestCase):
                 # Test not-owned SID
                 self.assertIsNone(txn)
 
+    def test_handle_sid_removed_from_universe(self):
+        # post some trades in the market
+        sim_params, _, _ = create_random_simulation_parameters()
+        events = factory.create_trade_history(
+            1,
+            [10, 10, 10, 10, 10],
+            [100, 100, 100, 100, 100],
+            oneday,
+            sim_params
+        )
+
+        # Create a tracker and a dividend
+        perf_tracker = perf.PerformanceTracker(sim_params)
+        dividend = factory.create_dividend(
+            1,
+            10.00,
+            # declared date, when the algorithm finds out about
+            # the dividend
+            events[0].dt,
+            # ex_date, the date before which the algorithm must hold stock
+            # to receive the dividend
+            events[1].dt,
+            # pay date, when the algorithm receives the dividend.
+            events[2].dt
+        )
+        dividend_frame = pd.DataFrame(
+            [dividend.to_series(index=zp.DIVIDEND_FIELDS)],
+        )
+        perf_tracker.update_dividends(dividend_frame)
+
+        # Ensure that the dividend is in the tracker
+        self.assertIn(1, perf_tracker.dividend_frame['sid'].values)
+
+        # Inform the tracker that sid 1 has been removed from the universe
+        perf_tracker.handle_sid_removed_from_universe(1)
+
+        # Ensure that the dividend for sid 1 has been removed from dividend
+        # frame
+        self.assertNotIn(1, perf_tracker.dividend_frame['sid'].values)
+
     def test_serialization(self):
         start_dt = datetime(year=2008,
                             month=10,
diff --git a/zipline/finance/performance/tracker.py b/zipline/finance/performance/tracker.py
index a55d92a1..e7beeb2e 100644
--- a/zipline/finance/performance/tracker.py
+++ b/zipline/finance/performance/tracker.py
@@ -236,6 +236,22 @@ class PerformanceTracker(object):
         self.dividend_frame = other.dividend_frame
         self._dividend_count = other._dividend_count
 
+    def handle_sid_removed_from_universe(self, sid):
+        """
+        This method handles any behaviors that must occur when a SID leaves the
+        universe of the TradingAlgorithm.
+
+        Parameters
+        __________
+        sid : int
+            The sid of the Asset being removed from the universe.
+        """
+
+        # Drop any dividends for the sid from the dividends frame
+        self.dividend_frame = self.dividend_frame[
+            self.dividend_frame.sid != sid
+        ]
+
     def update_performance(self):
         # calculate performance as of last trade
         for perf_period in self.perf_periods: