From cc58bb85cf6a3121b89c731391cd99cb80a3182d Mon Sep 17 00:00:00 2001
From: Kishan Patel <46427541+thekishanpatel@users.noreply.github.com>
Date: Mon, 9 Nov 2020 22:48:13 -0500
Subject: [PATCH] Checks if asset already in portfolio

Checks if asset already in portfolio in the add_position() method.
---
 pyrobot/portfolio.py | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

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]