diff --git a/pyrobot/portfolio.py b/pyrobot/portfolio.py index 1c81c10..26c9a53 100644 --- a/pyrobot/portfolio.py +++ b/pyrobot/portfolio.py @@ -148,18 +148,25 @@ def add_position(self, symbol: str, asset_type: str, purchase_date: Optional[str 'purchase_date': '2020-01-31' } """ - - self.positions[symbol] = {} - self.positions[symbol]['symbol'] = symbol - self.positions[symbol]['quantity'] = quantity - self.positions[symbol]['purchase_price'] = purchase_price - self.positions[symbol]['purchase_date'] = purchase_date - self.positions[symbol]['asset_type'] = asset_type - - if purchase_date: - self.positions[symbol]['ownership_status'] = True + if symbol not in self.positions: + self.positions[symbol] = {} + self.positions[symbol]['symbol'] = symbol + self.positions[symbol]['quantity'] = quantity + self.positions[symbol]['purchase_price'] = purchase_price + self.positions[symbol]['purchase_date'] = purchase_date + self.positions[symbol]['asset_type'] = asset_type + + if purchase_date: + self.positions[symbol]['ownership_status'] = True + else: + self.positions[symbol]['ownership_status'] = False else: - self.positions[symbol]['ownership_status'] = False + orig_price = self.positions[symbol].get('purchase_price', 0.00) + orig_quant = self.positions[symbol].get('quantity', 0) + avg_price = (orig_price*orig_quant + purchase_price*quantity) / (orig_quant + quantity) + new_quant = orig_quant + quantity + self.positions[symbol]['purchase_price'] = avg_price + self.positions[symbol]['quantity'] = new_quant return self.positions[symbol]