Total Complexity | 55 |
Total Lines | 385 |
Duplicated Lines | 0 % |
Complex classes like zipline.finance.performance.PositionTracker often do a lot of different things. To break such a class down, we need to identify a cohesive component within that class. A common approach to find such a component is to look for fields/methods that share the same prefixes, or suffixes.
Once you have determined the fields that belong together, you can apply the Extract Class refactoring. If the component makes sense as a sub-class, Extract Subclass is also a candidate, and is often faster.
1 | # |
||
124 | class PositionTracker(object): |
||
125 | |||
126 | def __init__(self, asset_finder, data_portal, data_frequency): |
||
127 | self.asset_finder = asset_finder |
||
128 | |||
129 | # FIXME really want to avoid storing a data portal here, |
||
130 | # but the path to get to maybe_create_close_position_transaction |
||
131 | # is long and tortuous |
||
132 | self._data_portal = data_portal |
||
133 | |||
134 | # sid => position object |
||
135 | self.positions = positiondict() |
||
136 | # Arrays for quick calculations of positions value |
||
137 | self._position_value_multipliers = OrderedDict() |
||
138 | self._position_exposure_multipliers = OrderedDict() |
||
139 | self._position_payout_multipliers = OrderedDict() |
||
140 | self._unpaid_dividends = {} |
||
141 | self._unpaid_stock_dividends = {} |
||
142 | self._positions_store = zp.Positions() |
||
143 | |||
144 | # Dict, keyed on dates, that contains lists of close position events |
||
145 | # for any Assets in this tracker's positions |
||
146 | self._auto_close_position_sids = {} |
||
147 | |||
148 | self.data_frequency = data_frequency |
||
149 | |||
150 | def _update_asset(self, sid): |
||
151 | try: |
||
152 | self._position_value_multipliers[sid] |
||
153 | self._position_exposure_multipliers[sid] |
||
154 | self._position_payout_multipliers[sid] |
||
155 | except KeyError: |
||
156 | # Check if there is an AssetFinder |
||
157 | if self.asset_finder is None: |
||
158 | raise PositionTrackerMissingAssetFinder() |
||
159 | |||
160 | # Collect the value multipliers from applicable sids |
||
161 | asset = self.asset_finder.retrieve_asset(sid) |
||
162 | if isinstance(asset, Equity): |
||
163 | self._position_value_multipliers[sid] = 1 |
||
164 | self._position_exposure_multipliers[sid] = 1 |
||
165 | self._position_payout_multipliers[sid] = 0 |
||
166 | if isinstance(asset, Future): |
||
167 | self._position_value_multipliers[sid] = 0 |
||
168 | self._position_exposure_multipliers[sid] = \ |
||
169 | asset.contract_multiplier |
||
170 | self._position_payout_multipliers[sid] = \ |
||
171 | asset.contract_multiplier |
||
172 | # Futures auto-close timing is controlled by the Future's |
||
173 | # auto_close_date property |
||
174 | self._insert_auto_close_position_date( |
||
175 | dt=asset.auto_close_date, |
||
176 | sid=sid |
||
177 | ) |
||
178 | |||
179 | def _insert_auto_close_position_date(self, dt, sid): |
||
180 | """ |
||
181 | Inserts the given SID in to the list of positions to be auto-closed by |
||
182 | the given dt. |
||
183 | |||
184 | Parameters |
||
185 | ---------- |
||
186 | dt : pandas.Timestamp |
||
187 | The date before-which the given SID will be auto-closed |
||
188 | sid : int |
||
189 | The SID of the Asset to be auto-closed |
||
190 | """ |
||
191 | if dt is not None: |
||
192 | self._auto_close_position_sids.setdefault(dt, set()).add(sid) |
||
193 | |||
194 | def auto_close_position_events(self, next_trading_day): |
||
195 | """ |
||
196 | Generates CLOSE_POSITION events for any SIDs whose auto-close date is |
||
197 | before or equal to the given date. |
||
198 | |||
199 | Parameters |
||
200 | ---------- |
||
201 | next_trading_day : pandas.Timestamp |
||
202 | The time before-which certain Assets need to be closed |
||
203 | |||
204 | Yields |
||
205 | ------ |
||
206 | Event |
||
207 | A close position event for any sids that should be closed before |
||
208 | the next_trading_day parameter |
||
209 | """ |
||
210 | past_asset_end_dates = set() |
||
211 | |||
212 | # Check the auto_close_position_dates dict for SIDs to close |
||
213 | for date, sids in self._auto_close_position_sids.items(): |
||
214 | if date > next_trading_day: |
||
215 | continue |
||
216 | past_asset_end_dates.add(date) |
||
217 | |||
218 | for sid in sids: |
||
219 | # Yield a CLOSE_POSITION event |
||
220 | event = Event({ |
||
221 | 'dt': date, |
||
222 | 'type': DATASOURCE_TYPE.CLOSE_POSITION, |
||
223 | 'sid': sid, |
||
224 | }) |
||
225 | yield event |
||
226 | |||
227 | # Clear out past dates |
||
228 | while past_asset_end_dates: |
||
229 | self._auto_close_position_sids.pop(past_asset_end_dates.pop()) |
||
230 | |||
231 | def update_positions(self, positions): |
||
232 | # update positions in batch |
||
233 | self.positions.update(positions) |
||
234 | for sid, pos in iteritems(positions): |
||
235 | self._update_asset(sid) |
||
236 | |||
237 | def update_position(self, sid, amount=None, last_sale_price=None, |
||
238 | last_sale_date=None, cost_basis=None): |
||
239 | if sid not in self.positions: |
||
240 | position = Position(sid) |
||
241 | self.positions[sid] = position |
||
242 | else: |
||
243 | position = self.positions[sid] |
||
244 | |||
245 | if amount is not None: |
||
246 | position.amount = amount |
||
247 | self._update_asset(sid=sid) |
||
248 | if last_sale_price is not None: |
||
249 | position.last_sale_price = last_sale_price |
||
250 | if last_sale_date is not None: |
||
251 | position.last_sale_date = last_sale_date |
||
252 | if cost_basis is not None: |
||
253 | position.cost_basis = cost_basis |
||
254 | |||
255 | def execute_transaction(self, txn): |
||
256 | # Update Position |
||
257 | # ---------------- |
||
258 | sid = txn.sid |
||
259 | |||
260 | if sid not in self.positions: |
||
261 | position = Position(sid) |
||
262 | self.positions[sid] = position |
||
263 | else: |
||
264 | position = self.positions[sid] |
||
265 | |||
266 | position.update(txn) |
||
267 | self._update_asset(sid) |
||
268 | |||
269 | def handle_commission(self, sid, cost): |
||
270 | # Adjust the cost basis of the stock if we own it |
||
271 | if sid in self.positions: |
||
272 | self.positions[sid].adjust_commission_cost_basis(sid, cost) |
||
273 | |||
274 | def handle_splits(self, splits): |
||
275 | """ |
||
276 | Processes a list of splits by modifying any positions as needed. |
||
277 | |||
278 | Parameters |
||
279 | ---------- |
||
280 | splits: list |
||
281 | A list of splits. Each split is a tuple of (sid, ratio). |
||
282 | |||
283 | Returns |
||
284 | ------- |
||
285 | None |
||
286 | """ |
||
287 | for split in splits: |
||
288 | sid = split[0] |
||
289 | if sid in self.positions: |
||
290 | # Make the position object handle the split. It returns the |
||
291 | # leftover cash from a fractional share, if there is any. |
||
292 | position = self.positions[sid] |
||
293 | leftover_cash = position.handle_split(sid, split[1]) |
||
294 | self._update_asset(split[0]) |
||
295 | return leftover_cash |
||
296 | |||
297 | def earn_dividends(self, dividends, stock_dividends): |
||
298 | """ |
||
299 | Given a list of dividends whose ex_dates are all the next trading day, |
||
300 | calculate and store the cash and/or stock payments to be paid on each |
||
301 | dividend's pay date. |
||
302 | """ |
||
303 | for dividend in dividends: |
||
304 | # Store the earned dividends so that they can be paid on the |
||
305 | # dividends' pay_dates. |
||
306 | div_owed = self.positions[dividend.sid].earn_dividend(dividend) |
||
307 | try: |
||
308 | self._unpaid_dividends[dividend.pay_date].append( |
||
309 | div_owed) |
||
310 | except KeyError: |
||
311 | self._unpaid_dividends[dividend.pay_date] = [div_owed] |
||
312 | |||
313 | for stock_dividend in stock_dividends: |
||
314 | div_owed = self.positions[stock_dividend.sid].earn_stock_dividend( |
||
315 | stock_dividend) |
||
316 | try: |
||
317 | self._unpaid_stock_dividends[stock_dividend.pay_date].\ |
||
318 | append(div_owed) |
||
319 | except KeyError: |
||
320 | self._unpaid_stock_dividends[stock_dividend.pay_date] = \ |
||
321 | [div_owed] |
||
322 | |||
323 | def pay_dividends(self, next_trading_day): |
||
324 | """ |
||
325 | Returns a cash payment based on the dividends that should be paid out |
||
326 | according to the accumulated bookkeeping of earned, unpaid, and stock |
||
327 | dividends. |
||
328 | """ |
||
329 | net_cash_payment = 0.0 |
||
330 | |||
331 | try: |
||
332 | payments = self._unpaid_dividends[next_trading_day] |
||
333 | # Mark these dividends as paid by dropping them from our unpaid |
||
334 | del self._unpaid_dividends[next_trading_day] |
||
335 | except KeyError: |
||
336 | payments = [] |
||
337 | |||
338 | # representing the fact that we're required to reimburse the owner of |
||
339 | # the stock for any dividends paid while borrowing. |
||
340 | for payment in payments: |
||
341 | net_cash_payment += payment['amount'] |
||
342 | |||
343 | # Add stock for any stock dividends paid. Again, the values here may |
||
344 | # be negative in the case of short positions. |
||
345 | |||
346 | try: |
||
347 | stock_payments = self._unpaid_stock_dividends[next_trading_day] |
||
348 | except: |
||
349 | stock_payments = [] |
||
350 | |||
351 | for stock_payment in stock_payments: |
||
352 | stock = stock_payment['payment_sid'] |
||
353 | share_count = stock_payment['share_count'] |
||
354 | # note we create a Position for stock dividend if we don't |
||
355 | # already own the asset |
||
356 | if stock in self.positions: |
||
357 | position = self.positions[stock] |
||
358 | else: |
||
359 | position = self.positions[stock] = Position(stock) |
||
360 | |||
361 | position.amount += share_count |
||
362 | self._update_asset(stock) |
||
363 | |||
364 | return net_cash_payment |
||
365 | |||
366 | def maybe_create_close_position_transaction(self, event): |
||
367 | if not self.positions.get(event.sid): |
||
368 | return None |
||
369 | |||
370 | amount = self.positions.get(event.sid).amount |
||
371 | price = self._data_portal.get_spot_value( |
||
372 | event.sid, 'close', event.dt, self.data_frequency) |
||
373 | |||
374 | txn = Transaction( |
||
375 | sid=event.sid, |
||
376 | amount=(-1 * amount), |
||
377 | dt=event.dt, |
||
378 | price=price, |
||
379 | commission=0, |
||
380 | order_id=0 |
||
381 | ) |
||
382 | return txn |
||
383 | |||
384 | def get_positions(self): |
||
385 | |||
386 | positions = self._positions_store |
||
387 | |||
388 | for sid, pos in iteritems(self.positions): |
||
389 | |||
390 | if pos.amount == 0: |
||
391 | # Clear out the position if it has become empty since the last |
||
392 | # time get_positions was called. Catching the KeyError is |
||
393 | # faster than checking `if sid in positions`, and this can be |
||
394 | # potentially called in a tight inner loop. |
||
395 | try: |
||
396 | del positions[sid] |
||
397 | except KeyError: |
||
398 | pass |
||
399 | continue |
||
400 | |||
401 | # Note that this will create a position if we don't currently have |
||
402 | # an entry |
||
403 | position = positions[sid] |
||
404 | position.amount = pos.amount |
||
405 | position.cost_basis = pos.cost_basis |
||
406 | position.last_sale_price = pos.last_sale_price |
||
407 | position.last_sale_date = pos.last_sale_date |
||
408 | |||
409 | return positions |
||
410 | |||
411 | def get_positions_list(self): |
||
412 | positions = [] |
||
413 | for sid, pos in iteritems(self.positions): |
||
414 | if pos.amount != 0: |
||
415 | positions.append(pos.to_dict()) |
||
416 | return positions |
||
417 | |||
418 | def sync_last_sale_prices(self, dt): |
||
419 | data_portal = self._data_portal |
||
420 | for sid, position in iteritems(self.positions): |
||
421 | position.last_sale_price = data_portal.get_spot_value( |
||
422 | sid, 'close', dt, self.data_frequency) |
||
423 | |||
424 | def stats(self): |
||
425 | amounts = [] |
||
426 | last_sale_prices = [] |
||
427 | for pos in itervalues(self.positions): |
||
428 | amounts.append(pos.amount) |
||
429 | last_sale_prices.append(pos.last_sale_price) |
||
430 | |||
431 | position_values = calc_position_values( |
||
432 | amounts, |
||
433 | last_sale_prices, |
||
434 | self._position_value_multipliers |
||
435 | ) |
||
436 | |||
437 | position_exposures = calc_position_exposures( |
||
438 | amounts, |
||
439 | last_sale_prices, |
||
440 | self._position_exposure_multipliers |
||
441 | ) |
||
442 | |||
443 | long_value = calc_long_value(position_values) |
||
444 | short_value = calc_short_value(position_values) |
||
445 | gross_value = calc_gross_value(long_value, short_value) |
||
446 | long_exposure = calc_long_exposure(position_exposures) |
||
447 | short_exposure = calc_short_exposure(position_exposures) |
||
448 | gross_exposure = calc_gross_exposure(long_exposure, short_exposure) |
||
449 | net_exposure = calc_net(position_exposures) |
||
450 | longs_count = calc_longs_count(position_exposures) |
||
451 | shorts_count = calc_shorts_count(position_exposures) |
||
452 | net_value = calc_net(position_values) |
||
453 | |||
454 | return PositionStats( |
||
455 | long_value=long_value, |
||
456 | gross_value=gross_value, |
||
457 | short_value=short_value, |
||
458 | long_exposure=long_exposure, |
||
459 | short_exposure=short_exposure, |
||
460 | gross_exposure=gross_exposure, |
||
461 | net_exposure=net_exposure, |
||
462 | longs_count=longs_count, |
||
463 | shorts_count=shorts_count, |
||
464 | net_value=net_value |
||
465 | ) |
||
466 | |||
467 | def __getstate__(self): |
||
468 | state_dict = {} |
||
469 | |||
470 | state_dict['asset_finder'] = self.asset_finder |
||
471 | state_dict['positions'] = dict(self.positions) |
||
472 | state_dict['unpaid_dividends'] = self._unpaid_dividends |
||
473 | state_dict['unpaid_stock_dividends'] = self._unpaid_stock_dividends |
||
474 | state_dict['auto_close_position_sids'] = self._auto_close_position_sids |
||
475 | state_dict['data_frequency'] = self.data_frequency |
||
476 | |||
477 | STATE_VERSION = 3 |
||
478 | state_dict[VERSION_LABEL] = STATE_VERSION |
||
479 | return state_dict |
||
480 | |||
481 | def __setstate__(self, state): |
||
482 | OLDEST_SUPPORTED_STATE = 3 |
||
483 | version = state.pop(VERSION_LABEL) |
||
484 | |||
485 | if version < OLDEST_SUPPORTED_STATE: |
||
486 | raise BaseException("PositionTracker saved state is too old.") |
||
487 | |||
488 | self.asset_finder = state['asset_finder'] |
||
489 | self.positions = positiondict() |
||
490 | self.data_frequency = state['data_frequency'] |
||
491 | # note that positions_store is temporary and gets regened from |
||
492 | # .positions |
||
493 | self._positions_store = zp.Positions() |
||
494 | |||
495 | self._unpaid_dividends = state['unpaid_dividends'] |
||
496 | self._unpaid_stock_dividends = state['unpaid_stock_dividends'] |
||
497 | self._auto_close_position_sids = state['auto_close_position_sids'] |
||
498 | |||
499 | # Arrays for quick calculations of positions value |
||
500 | self._position_value_multipliers = OrderedDict() |
||
501 | self._position_exposure_multipliers = OrderedDict() |
||
502 | self._position_payout_multipliers = OrderedDict() |
||
503 | |||
504 | # Update positions is called without a finder |
||
505 | self.update_positions(state['positions']) |
||
506 | |||
507 | # FIXME |
||
508 | self._data_portal = None |
||
509 |