From 708c68b3cbd09d9bb1a99291044bbc19198b8499 Mon Sep 17 00:00:00 2001
From: Shrivu Shankar <shrivu1122@gmail.com>
Date: Mon, 20 Aug 2018 13:00:26 -0500
Subject: [PATCH] Added get_assets() fixed #2

---
 README.md               |   2 +
 Tradinhood/Robinhood.py | 107 ++++++++++++++++++++++++++--------------
 2 files changed, 72 insertions(+), 37 deletions(-)

diff --git a/README.md b/README.md
index 079d00a..4af00c0 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,8 @@ for order in rbh.orders:
     if order.state == 'confirmed' and order.asset_type == 'cryptocurrency':
         order.cancel()
         rbh.wait_for_orders([order]) # wait for cancel
+
+rbh.get_assets() # get everything you own as a dict(asset:amt)
 ```
 
 ## Dataset
diff --git a/Tradinhood/Robinhood.py b/Tradinhood/Robinhood.py
index 73452ae..63bc4cf 100644
--- a/Tradinhood/Robinhood.py
+++ b/Tradinhood/Robinhood.py
@@ -200,50 +200,19 @@ def quantity(self, asset, include_held=False):
         """
         assert self.logged_in
 
-        if isinstance(asset, str):
+        if isinstance(asset, str): # convert str to Stock or Currency
             asset = self.__getitem__(asset)
 
         if isinstance(asset, Currency):
-
-            currs = self.holdings
-
-            for curr in currs:
-
-                if curr['currency']['code'] == asset.code:
-
-                    amt = Decimal(curr['quantity_available'])
-
-                    if include_held:
-                        amt += Decimal(curr['quantity_held_for_buy'])
-                        amt += Decimal(curr['quantity_held_for_sell'])
-
-                    return amt
-
-            return Decimal('0.00')
+            assets = self.get_assets(include_positions=False, include_holdings=True, include_held=include_held)
 
         elif isinstance(asset, Stock):
-
-            stocks = self.positions
-
-            for stock in stocks:
-
-                if asset.id in stock['instrument']:
-
-                    amt = Decimal(stock['quantity'])
-
-                    if include_held:
-                        amt += Decimal(stock['shares_held_for_buys'])
-                        amt += Decimal(stock['shares_held_for_sells'])
-                        amt += Decimal(stock['shares_held_for_options_collateral'])
-                        amt += Decimal(stock['shares_held_for_options_events'])
-                        amt += Decimal(stock['shares_held_for_stock_grants'])
-
-                    return amt
-
-            return Decimal('0.00')
+            assets = self.get_assets(include_positions=True, include_holdings=False, include_held=include_held)
 
         else:
-            raise UsageError('Invalid asset')
+            raise UsageError('Invalid asset type')
+
+        return assets.get(asset, Decimal('0.00')) # default to zero if not in positions or holdings
 
     def _order(self, order_side, asset, amt, type='market', price=None, stop_price=None, time_in_force='gtc', return_json=False):
         """Internal order method
@@ -438,6 +407,61 @@ def wait_for_orders(self, orders, delay=5, timeout=120, force=False):
 
         return all(map(order_complete, orders))
 
+    def get_assets(self, include_positions=True, include_holdings=True, include_held=False):
+        """Get all owned assets
+
+        Args:
+            include_positions: (bool) whether to include stocks
+            include_holdings: (bool) whether to include currencies
+            include_held: (bool) whether to include held assets
+
+        Returns:
+            (dict) Stock or Currency objects paired with quantities, note:
+                some quantities may be zero
+        """
+        assert self.logged_in
+
+        my_assets = {}
+
+        if include_positions:
+
+            stocks = self.positions
+
+            for stock_json in stocks:
+
+                stock = Stock.from_url(self.session, stock_json['instrument'])
+                amt = Decimal(stock_json['quantity'])
+
+                if include_held:
+                    amt += Decimal(stock_json['shares_held_for_buys'])
+                    amt += Decimal(stock_json['shares_held_for_sells'])
+                    amt += Decimal(stock_json['shares_held_for_options_collateral'])
+                    amt += Decimal(stock_json['shares_held_for_options_events'])
+                    amt += Decimal(stock_json['shares_held_for_stock_grants'])
+
+                my_assets[stock] = amt
+
+        if include_holdings:
+
+            currs = self.holdings
+
+            for curr_json in currs:
+
+                code = curr_json['currency']['code']
+
+                if code in Currency.cache: # all currencies already cached
+
+                    curr = Currency.cache[code]
+                    amt = Decimal(curr_json['quantity_available'])
+
+                    if include_held:
+                        amt += Decimal(curr_json['quantity_held_for_buy'])
+                        amt += Decimal(curr_json['quantity_held_for_sell'])
+
+                    my_assets[curr] = amt
+
+        return my_assets
+
     @property
     def account_info(self):
         """Account info"""
@@ -595,6 +619,15 @@ def __init__(self, session, instrument_json):
 
         Stock.cache[self.symbol] = self
 
+    @classmethod
+    def from_url(self, session, instrument_url):
+        """Create a stock from its instrument url"""
+        for symbol, stock in Stock.cache.items(): # try cache
+            if stock.id in instrument_url:
+                return stock
+
+        return Stock(session, session.get(instrument_url).json())
+
     @property
     def market_open(self):
         """If the market for this stock is open"""