Refactor disc commands to be more efficient

gamestonk_terminal/discovery/disc_controller.py

@@ -42,22 +42,14 @@ class DiscoveryController:
         "simply_wallst",
         "spachero",
         "uwhales",
-        "sec_val",
-        "sec_perf",
-        "sec_spec",
-        "ind_val",
-        "ind_perf",
-        "ind_spec",
-        "cntry_val",
-        "cntry_perf",
-        "cntry_spec",
+        "valuation",
+        "performance",
+        "spectrum",
     ]
 
     def __init__(self):
         """Constructor"""
-        self.delete_sector_spectrum = False
-        self.delete_industry_spectrum = False
-        self.delete_country_spectrum = False
+        self.spectrum_img_to_delete = ""
         self.disc_parser = argparse.ArgumentParser(add_help=False, prog="disc")
         self.disc_parser.add_argument(
             "cmd",
@@ -91,17 +83,9 @@ class DiscoveryController:
         print("   simply_wallst  Simply Wall St. research data [Simply Wall St.]")
         print("   spachero       great website for SPACs research [SpacHero]")
         print("   uwhales        good website for SPACs research [UnusualWhales]")
-        print("")
-        print("Finviz:")
-        print("   sec_val        sector valuation")
-        print("   sec_perf       sector performance")
-        print("   sec_spec       sector spectrum")
-        print("   ind_val        industry valuation")
-        print("   ind_perf       industry performance")
-        print("   ind_spec       industry spectrum")
-        print("   cntry_val      country valuation")
-        print("   cntry_perf     country performance")
-        print("   cntry_spec     country spectrum")
+        print("   valuation      valuation of sectors, industry, country [Finviz]")
+        print("   performance    performance of sectors, industry, country [Finviz]")
+        print("   spectrum       spectrum of sectors, industry, country [Finviz]")
         print("")
 
     def switch(self, an_input: str):
@@ -118,15 +102,9 @@ class DiscoveryController:
 
         # Due to Finviz implementation of Spectrum, we delete the generated spectrum figure
         # after saving it and displaying it to the user
-        if self.delete_sector_spectrum:
-            os.remove("Sector.jpg")
-            self.delete_sector_spectrum = False
-        elif self.delete_industry_spectrum:
-            os.remove("Industry.jpg")
-            self.delete_industry_spectrum = False
-        elif self.delete_country_spectrum:
-            os.remove("Country (U.S. listed stocks only).jpg")
-            self.delete_country_spectrum = False
+        if self.spectrum_img_to_delete:
+            os.remove(self.spectrum_img_to_delete + ".jpg")
+            self.spectrum_img_to_delete = ""
 
         return getattr(
             self, "call_" + known_args.cmd, lambda: "Command not recognized!"
@@ -192,51 +170,20 @@ class DiscoveryController:
         """Process uwhales command"""
         unusual_whales_view.unusual_whales_view(other_args)
 
-    def call_sec_val(self, other_args: List[str]):
-        """Process sec_val command"""
-        finviz_view.view_group_data(other_args, "Sector", "valuation")
+    def call_valuation(self, other_args: List[str]):
+        """Process valuation command"""
+        finviz_view.view_group_data(other_args, "valuation")
 
-    def call_sec_perf(self, other_args: List[str]):
-        """Process sec_perf command"""
-        finviz_view.view_group_data(other_args, "Sector", "performance")
+    def call_performance(self, other_args: List[str]):
+        """Process performance command"""
+        finviz_view.view_group_data(other_args, "performance")
 
-    def call_sec_spec(self, other_args: List[str]):
-        """Process sec_spec command"""
-        finviz_view.view_group_data(other_args, "Sector", "spectrum")
-        self.delete_sector_spectrum = True
-
-    def call_ind_val(self, other_args: List[str]):
-        """Process ind_val command"""
-        finviz_view.view_group_data(other_args, "Industry", "valuation")
-
-    def call_ind_perf(self, other_args: List[str]):
-        """Process ind_perf command"""
-        finviz_view.view_group_data(other_args, "Industry", "performance")
-
-    def call_ind_spec(self, other_args: List[str]):
-        """Process ind_spec command"""
-        finviz_view.view_group_data(other_args, "Industry", "spectrum")
-        self.delete_industry_spectrum = True
-
-    def call_cntry_val(self, other_args: List[str]):
-        """Process cntry_val command"""
-        finviz_view.view_group_data(
-            other_args, "Country (U.S. listed stocks only)", "valuation"
+    def call_spectrum(self, other_args: List[str]):
+        """Process spectrum command"""
+        self.spectrum_img_to_delete = finviz_view.view_group_data(
+            other_args, "spectrum"
         )
 
-    def call_cntry_perf(self, other_args: List[str]):
-        """Process cntry_perf command"""
-        finviz_view.view_group_data(
-            other_args, "Country (U.S. listed stocks only)", "performance"
-        )
-
-    def call_cntry_spec(self, other_args: List[str]):
-        """Process cntry_spec command"""
-        finviz_view.view_group_data(
-            other_args, "Country (U.S. listed stocks only)", "spectrum"
-        )
-        self.delete_country_spectrum = True
-
 
 def menu():
     """Discovery Menu"""

gamestonk_terminal/discovery/finviz_view.py

@@ -8,6 +8,24 @@ from gamestonk_terminal.helper_funcs import parse_known_args_and_warn
 from finvizfinance.group import valuation, performance, spectrum
 from PIL import Image
 
+l_GROUPS = [
+    "Sector",
+    "Industry",
+    "Industry (Basic Materials)",
+    "Industry (Communication Services)",
+    "Industry (Consumer Cyclical)",
+    "Industry (Consumer Defensive)",
+    "Industry (Energy)",
+    "Industry (Financial)",
+    "Industry (Healthcare)",
+    "Industry (Industrials)",
+    "Industry (Real Estate)",
+    "Industry (Technology)",
+    "Industry (Utilities)",
+    "Country (U.S. listed stocks only)",
+    "Capitalization",
+]
+
 
 def map_sp500_view(other_args: List[str]):
     """Opens Finviz website in a browser
@@ -102,15 +120,13 @@ def get_spectrum_data(group: str):
     spectrum.Spectrum().ScreenerView(group=group)
 
 
-def view_group_data(other_args: List[str], group: str, data_type: str):
+def view_group_data(other_args: List[str], data_type: str):
     """View group (sectors, industry or country) valuation/performance/spectrum data
 
     Parameters
     ----------
     other_args : List[str]
         argparse other args
-    group : str
-        select group to see data between sectors, industry and country
     data_type : str
         select data type to see data between valuation, performance and spectrum
     """
@@ -121,24 +137,50 @@ def view_group_data(other_args: List[str], group: str, data_type: str):
             View group (sectors, industry or country) valuation/performance/spectrum data. [Source: Finviz]
         """,
     )
-    ns_parser = parse_known_args_and_warn(parser, other_args)
-    if not ns_parser:
-        return
+    parser.add_argument(
+        "-g",
+        "--group",
+        nargs="+",
+        type=str,
+        default="Sector",
+        dest="group",
+        help="Data group (sector, industry or country)",
+        choices=l_GROUPS,
+    )
 
-    if data_type == "valuation":
-        df_group = get_valuation_performance_data(group, "valuation")
-        print(df_group.to_string())
+    try:
+        if other_args:
+            if "-" not in other_args[0]:
+                other_args.insert(0, "-g")
+            other_args = [other_args[0], " ".join(other_args[1:])]
 
-    elif data_type == "performance":
-        df_group = get_valuation_performance_data(group, "performance")
-        print(df_group.to_string())
+        ns_parser = parse_known_args_and_warn(parser, other_args)
+        if not ns_parser:
+            return
 
-    elif data_type == "spectrum":
-        get_spectrum_data(group)
+        if isinstance(ns_parser.group, List):
+            ns_parser.group = ns_parser.group[0]
 
-        img = Image.open(group + ".jpg")
-        img.show()
-    else:
-        print("Invalid data type. Choose between valuation, performance and spectrum.")
+        if data_type == "valuation" or data_type == "performance":
+            df_group = get_valuation_performance_data(ns_parser.group, data_type)
+            print(df_group.to_string())
 
-    print("")
+        elif data_type == "spectrum":
+            get_spectrum_data(ns_parser.group)
+
+            img = Image.open(ns_parser.group + ".jpg")
+            img.show()
+
+            print("")
+            return ns_parser.group
+
+        else:
+            print(
+                "Invalid data type. Choose between valuation, performance and spectrum."
+            )
+
+        print("")
+        return ""
+
+    except SystemExit:
+        print("")