mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-10-27 12:40:35 +00:00
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
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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(" ", "_"))
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user