Skip to content

Commit

Permalink
ENH: Adds a perf tracker method to handle SIDs leaving the universe
Browse files Browse the repository at this point in the history
  • Loading branch information
jfkirk committed Jul 9, 2015
1 parent 7cde393 commit efa6d8d
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 0 deletions.
40 changes: 40 additions & 0 deletions tests/test_perf_tracking.py
Expand Up @@ -2020,6 +2020,46 @@ def test_close_position_event(self, env=None):
# 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,
Expand Down
16 changes: 16 additions & 0 deletions zipline/finance/performance/tracker.py
Expand Up @@ -236,6 +236,22 @@ def initialize_dividends_from_other(self, other):
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:
Expand Down

0 comments on commit efa6d8d

Please sign in to comment.