added select/mktcap options

gamestonk_terminal/portfolio_optimization/port_opt_api.py

@@ -3,37 +3,9 @@ __docformat__ = "numpy"
 
 import argparse
 from typing import List
+import yfinance as yf
 
-def optimize(list_of_stocks:List[str], other_args:List[str]):
-    """
-
-    Parameters
-    ----------
-    list_of_stocks: List[str]
-        List of tickers to be included in optimization
-
-    Returns
-    -------
-    weights : dict
-        Dictionary of weights where keys are the tickers
-
-    """
-
-    parser = argparse.ArgumentParser(
-        add_help=False,
-        prog="optimize",
-    )
-
-    parse
-
-    weights = {}
-    n_stocks = len(list_of_stocks)
-    for stock in list_of_stocks:
-        weights[stock] = round(1/n_stocks, 5)
-
-    return weights
-
-def equal_weight(list_of_stocks:List[str]):
+def equal_weight(list_of_stocks:List[str], other_args:List[str]):
     """
         Equally weighted portfolio, where weight = 1/# of stocks
 
@@ -55,3 +27,27 @@ def equal_weight(list_of_stocks:List[str]):
 
     return weights
 
+def market_cap_weighting(list_of_stocks:List[str], other_args:List[str]):
+    """
+    Market cap weighted portfolio where each weight is the stocks market cap/ sum of all market caps
+    Parameters
+    ----------
+    list_of_stocks: List[str[
+        List of tickers to be included in optimization
+
+    Returns
+    -------
+    weights: dict
+        Dictionary of weights where keys are the tickers
+    """
+    weights = {}
+    mkt_cap = {}
+    total_cap = 0
+    for stock in list_of_stocks:
+        stock_cap = yf.Ticker(stock).info["marketCap"]
+        mkt_cap[stock] = stock_cap
+        total_cap += stock_cap
+    for k,v in mkt_cap.items():
+        weights[k] = round(v/total_cap,5)
+
+    return weights

gamestonk_terminal/portfolio_optimization/port_optimization_controller.py

@@ -86,11 +86,7 @@ class PortfolioOptimization:
     def call_add(self, other_args:List[str]):
         self.add_stocks(self, other_args)
 
-    def call_optimize(self, other_args:List[str]):
-        weights = po_api.optimize(self.tickers, other_args)
-        print("Optimal Weights for Equal Weighting:")
-        print(weights)
-        print("")
+
     def call_equal_weight(self, other_args:List[str]):
         weights = po_api.equal_weight(self.tickers, other_args)
         print("Optimal Weights for Equal Weighting:")
@@ -98,9 +94,14 @@ class PortfolioOptimization:
         print("")
 
     def call_mkt_cap(self,other_args:List[str]):
-        print("TODO")
+        weights = po_api.market_cap_weighting(self.tickers, other_args)
+        print("Market Cap Weighting Weights:")
+        print(weights)
+        print("")
+
     def call_select(self,other_args:List[str]):
-        print("TODO")
+        self.tickers = set([])
+        self.add_stocks(self, other_args)
 
     @staticmethod
     def add_stocks(self, other_args: List[str]):