From f9e6b20b60da9162d653a131a24d7caedfc7f5ba Mon Sep 17 00:00:00 2001
From: James Maslek <jmaslek11@gmail.com>
Date: Mon, 3 Apr 2023 06:38:02 -0400
Subject: [PATCH] Add --sheet-name to load on econometrics and forecast (#4525)

* Add sheet-name to econometrics load

* Also do forecasting load

* Handle error when sheet not found
---
 openbb_terminal/common/common_model.py          | 11 ++++++++++-
 .../econometrics/econometrics_controller.py     | 17 ++++++++++++++++-
 openbb_terminal/forecast/forecast_controller.py | 17 ++++++++++++++++-
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/openbb_terminal/common/common_model.py b/openbb_terminal/common/common_model.py
index 1130469b4de..bbf1d98cabf 100644
--- a/openbb_terminal/common/common_model.py
+++ b/openbb_terminal/common/common_model.py
@@ -52,6 +52,7 @@ file_types = ["xlsx", "csv"]
 @log_start_end(log=logger)
 def load(
     file: str,
+    sheet_name: Optional[str] = None,
     data_files: Optional[Dict[Any, Any]] = None,
     data_examples: Optional[Dict[Any, Any]] = None,
 ) -> pd.DataFrame:
@@ -90,7 +91,15 @@ def load(
 
     try:
         if file_type == ".xlsx":
-            data = pd.read_excel(full_file)
+            try:
+                data = (
+                    pd.read_excel(full_file)
+                    if sheet_name is None
+                    else pd.read_excel(full_file, sheet_name=sheet_name)
+                )
+            except ValueError as err:
+                console.print(f"[red]{err}[/red]\n")
+                return pd.DataFrame()
         elif file_type == ".csv":
             data = pd.read_csv(full_file)
         else:
diff --git a/openbb_terminal/econometrics/econometrics_controller.py b/openbb_terminal/econometrics/econometrics_controller.py
index 1e39c7d83b4..cf1233a1cf1 100644
--- a/openbb_terminal/econometrics/econometrics_controller.py
+++ b/openbb_terminal/econometrics/econometrics_controller.py
@@ -170,6 +170,7 @@ class EconometricsController(BaseController):
                 "-a": "-alias",
                 "--examples": None,
                 "-e": "--examples",
+                "--sheet-name": None,
             }
 
             for feature in ["export", "show", "desc", "clear", "index"]:
@@ -349,6 +350,13 @@ class EconometricsController(BaseController):
             default=False,
             dest="examples",
         )
+        parser.add_argument(
+            "--sheet-name",
+            dest="sheet_name",
+            default=None,
+            nargs="+",
+            help="Name of excel sheet to save data to. Only valid for .xlsx files.",
+        )
 
         if other_args and "-" not in other_args[0][0]:
             other_args.insert(0, "-f")
@@ -405,7 +413,14 @@ class EconometricsController(BaseController):
                 )
                 return
 
-            data = common_model.load(file, self.DATA_FILES, common_model.DATA_EXAMPLES)
+            data = common_model.load(
+                file,
+                data_files=self.DATA_FILES,
+                data_examples=common_model.DATA_EXAMPLES,
+                sheet_name=" ".join(ns_parser.sheet_name)
+                if ns_parser.sheet_name
+                else None,
+            )
 
             if not data.empty:
                 data.columns = data.columns.map(lambda x: x.lower().replace(" ", "_"))
diff --git a/openbb_terminal/forecast/forecast_controller.py b/openbb_terminal/forecast/forecast_controller.py
index 25bffa9ef78..467194b61e6 100644
--- a/openbb_terminal/forecast/forecast_controller.py
+++ b/openbb_terminal/forecast/forecast_controller.py
@@ -772,6 +772,14 @@ class ForecastController(BaseController):
             help="Alias name to give to the dataset",
             type=str,
         )
+        parser.add_argument(
+            "--sheet-name",
+            dest="sheet_name",
+            default=None,
+            nargs="+",
+            help="Name of excel sheet to save data to. Only valid for .xlsx files.",
+        )
+
         # Load in any newly exported files
         self.DATA_FILES = forecast_model.get_default_files()
 
@@ -796,7 +804,14 @@ class ForecastController(BaseController):
                         "[red]The file/dataset selected has already been loaded.[/red]\n"
                     )
                     return
-                data = common_model.load(file, self.DATA_FILES, {})
+                data = common_model.load(
+                    file,
+                    data_files=self.DATA_FILES,
+                    data_examples={},
+                    sheet_name=" ".join(ns_parser.sheet_name)
+                    if ns_parser.sheet_name
+                    else None,
+                )
                 if not data.empty:
                     self.files_full.append([ns_parser.file, ns_parser.alias])
                     self.load(alias, data)