[Feature] Improvements to handling charts on the CLI `results` (#6544)

* better handling of the chart arg: falls back to charting.to_chart() and only then, if not possible to display falls back to the print_rich_table()

* handling chart arguments aka unknown arguments

* propagating changes on return type of parse_simple_args to choices.py for coherency

5 files changed, 134 insertions, 47 deletions

diff --git a/cli/openbb_cli/controllers/base_controller.py b/cli/openbb_cli/controllers/base_controller.py
index 91a353f8b60..da54491c11e 100644
--- a/cli/openbb_cli/controllers/base_controller.py
+++ b/cli/openbb_cli/controllers/base_controller.py
@@ -8,7 +8,7 @@ import shlex
 from abc import ABCMeta, abstractmethod
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 import pandas as pd
 from openbb_cli.config.completer import NestedCompleter
@@ -19,7 +19,9 @@ from openbb_cli.controllers.utils import (
     check_file_type_saved,
     check_positive,
     get_flair_and_username,
+    handle_obbject_display,
     parse_and_split_input,
+    parse_unknown_args_to_dict,
     print_guest_block_msg,
     print_rich_table,
     remove_file,
@@ -373,7 +375,7 @@ class BaseController(metaclass=ABCMeta):
         if other_args and "-" not in other_args[0][0]:
             other_args.insert(0, "-n")
 
-        ns_parser = self.parse_simple_args(parser, other_args)
+        ns_parser, _ = self.parse_simple_args(parser, other_args)
 
         if ns_parser:
             if not ns_parser.name:
@@ -472,7 +474,7 @@ class BaseController(metaclass=ABCMeta):
             description="Stop recording session into .openbb routine file",
         )
         # This is only for auto-completion purposes
-        _ = self.parse_simple_args(parser, other_args)
+        _, _ = self.parse_simple_args(parser, other_args)
 
         if "-h" not in other_args and "--help" not in other_args:
             global RECORD_SESSION  # noqa: PLW0603
@@ -603,7 +605,7 @@ class BaseController(metaclass=ABCMeta):
             prog="whoami",
             description="Show current user",
         )
-        ns_parser = self.parse_simple_args(parser, other_args)
+        ns_parser, _ = self.parse_simple_args(parser, other_args)
 
         if ns_parser:
             current_user = session.user
@@ -635,11 +637,27 @@ class BaseController(metaclass=ABCMeta):
             "--chart", action="store_true", dest="chart", help="Display chart."
         )
         parser.add_argument(
-            "--export", dest="export", help="Export data.", nargs="+", default=None
+            "--export",
+            default="",
+            type=check_file_type_saved(["csv", "json", "xlsx", "png", "jpg"]),
+            dest="export",
+            help="Export raw data into csv, json, xlsx and figure into png or jpg.",
+            nargs="+",
+        )
+        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.",
+        )
+
+        ns_parser, unknown_args = self.parse_simple_args(
+            parser, other_args, unknown_args=True
         )
 
-        ns_parser = self.parse_simple_args(parser, other_args)
         if ns_parser:
+            kwargs = parse_unknown_args_to_dict(unknown_args)
             if not ns_parser.index and not ns_parser.key:
                 results = session.obbject_registry.all
                 if results:
@@ -657,21 +675,13 @@ class BaseController(metaclass=ABCMeta):
                     index = int(ns_parser.index)
                     obbject = session.obbject_registry.get(index)
                     if obbject:
-                        if ns_parser.chart and obbject.chart:
-                            obbject.show()
-                        else:
-                            title = obbject.extra.get("command", "")
-                            df = obbject.to_dataframe()
-                            print_rich_table(
-                                df=df,
-                                show_index=True,
-                                title=title,
-                                export=ns_parser.export,
-                            )
-                            if ns_parser.chart and not obbject.chart:
-                                session.console.print(
-                                    "[info]No chart available.[/info]"
-                                )
+                        handle_obbject_display(
+                            obbject=obbject,
+                            chart=ns_parser.chart,
+                            export=ns_parser.export,
+                            sheet_name=ns_parser.sheet_name,
+                            **kwargs,
+                        )
                     else:
                         session.console.print(
                             f"[info]No result found at index {index}.[/info]"
@@ -683,26 +693,24 @@ class BaseController(metaclass=ABCMeta):
             elif ns_parser.key:
                 obbject = session.obbject_registry.get(ns_parser.key)
                 if obbject:
-                    if ns_parser.chart and obbject.chart:
-                        obbject.show()
-                    else:
-                        title = obbject.extra.get("command", "")
-                        df = obbject.to_dataframe()
-                        print_rich_table(
-                            df=df,
-                            show_index=True,
-                            title=title,
-                            export=ns_parser.export,
-                        )
-                        if ns_parser.chart and not obbject.chart:
-                            session.console.print("[info]No chart available.[/info]")
+                    handle_obbject_display(
+                        obbject=obbject,
+                        chart=ns_parser.chart,
+                        export=ns_parser.export,
+                        sheet_name=ns_parser.sheet_name,
+                        **kwargs,
+                    )
                 else:
                     session.console.print(
                         f"[info]No result found with key '{ns_parser.key}'.[/info]"
                     )
 
     @staticmethod
-    def parse_simple_args(parser: argparse.ArgumentParser, other_args: List[str]):
+    def parse_simple_args(
+        parser: argparse.ArgumentParser,
+        other_args: List[str],
+        unknown_args: bool = False,
+    ) -> Tuple[Optional[argparse.Namespace], Optional[List[str]]]:
         """Parse list of arguments into the supplied parser.
 
         Parameters
@@ -711,11 +719,15 @@ class BaseController(metaclass=ABCMeta):
             Parser with predefined arguments
         other_args: List[str]
             List of arguments to parse
+        unknown_args: bool
+            Flag to indicate if unknown arguments should be returned
 
         Returns
         -------
-        ns_parser:
+        ns_parser: argparse.Namespace
             Namespace with parsed arguments
+        l_unknown_args: List[str]
+            List of unknown arguments
         """
         parser.add_argument(
             "-h", "--help", action="store_true", help="show this help message"
@@ -729,19 +741,18 @@ class BaseController(metaclass=ABCMeta):
         except SystemExit:
             # In case the command has required argument that isn't specified
             session.console.print("\n")
-            return None
+            return None, None
 
         if ns_parser.help:
             txt_help = parser.format_help()
             session.console.print(f"[help]{txt_help}[/help]")
-            return None
+            return None, None
 
-        if l_unknown_args:
+        if l_unknown_args and not unknown_args:
             session.console.print(
                 f"The following args couldn't be interpreted: {l_unknown_args}\n"
             )
-
-        return ns_parser
+        return ns_parser, l_unknown_args
 
     @classmethod
     def parse_known_args_and_warn(
diff --git a/cli/openbb_cli/controllers/choices.py b/cli/openbb_cli/controllers/choices.py
index 3afc5234c11..f4e87552095 100644
--- a/cli/openbb_cli/controllers/choices.py
+++ b/cli/openbb_cli/controllers/choices.py
@@ -4,7 +4,7 @@ from argparse import SUPPRESS, ArgumentParser
 from contextlib import contextmanager
 from inspect import isfunction, unwrap
 from types import MethodType
-from typing import Callable, List, Literal
+from typing import Callable, List, Literal, Tuple
 from unittest.mock import patch
 
 from openbb_cli.controllers.utils import (
@@ -110,7 +110,7 @@ def __mock_parse_known_args_and_warn(
     )
 
 
-def __mock_parse_simple_args(parser: ArgumentParser, other_args: List[str]) -> None:
+def __mock_parse_simple_args(parser: ArgumentParser, other_args: List[str]) -> Tuple:
     """Add arguments.
 
     Add the arguments that would have normally added by:
@@ -127,6 +127,7 @@ def __mock_parse_simple_args(parser: ArgumentParser, other_args: List[str]) -> N
         "-h", "--help", action="store_true", help="show this help message"
     )
     _ = other_args
+    return None, None
 
 
 def __get_command_func(controller, command: str):
@@ -216,7 +217,7 @@ def __patch_controller_functions(controller):
             target=controller,
             attribute="parse_simple_args",
             side_effect=__mock_parse_simple_args,
-            return_value=None,
+            return_value=(None, None),
         ),
         patch.object(
             target=controller,
diff --git a/cli/openbb_cli/controllers/cli_controller.py b/cli/openbb_cli/controllers/cli_controller.py
index 004239bf436..9bacd479505 100644
--- a/cli/openbb_cli/controllers/cli_controller.py
+++ b/cli/openbb_cli/controllers/cli_controller.py
@@ -205,10 +205,11 @@ class CLIController(BaseController):
             choices["results"] = {
                 "--help": None,
                 "-h": "--help",
-                "--export": None,
+                "--export": {c: None for c in ["csv", "json", "xlsx", "png", "jpg"]},
                 "--index": None,
                 "--key": None,
                 "--chart": None,
+                "--sheet_name": None,
             }
 
             self.update_completer(choices)
diff --git a/cli/openbb_cli/controllers/settings_controller.py b/cli/openbb_cli/controllers/settings_controller.py
index da8add50b52..ab1099362f0 100644
--- a/cli/openbb_cli/controllers/settings_controller.py
+++ b/cli/openbb_cli/controllers/settings_controller.py
@@ -80,7 +80,7 @@ class SettingsController(BaseController):
                 description=field["description"],
                 add_help=False,
             )
-            ns_parser = self.parse_simple_args(parser, other_args)
+            ns_parser, _ = self.parse_simple_args(parser, other_args)
             if ns_parser:
                 session.settings.set_item(
                     field_name, not getattr(session.settings, field_name)
@@ -113,7 +113,7 @@ class SettingsController(BaseController):
                 type=type_,  # type: ignore[arg-type]
                 choices=choices,
             )
-            ns_parser = self.parse_simple_args(parser, other_args)
+            ns_parser, _ = self.parse_simple_args(parser, other_args)
             if ns_parser:
                 if ns_parser.value:
                     # Console style is applied immediately
diff --git a/cli/openbb_cli/controllers/utils.py b/cli/openbb_cli/controllers/utils.py
index c921ad5ab3a..5ab5a817a5e 100644
--- a/cli/openbb_cli/controllers/utils.py
+++ b/cli/openbb_cli/controllers/utils.py
@@ -21,6 +21,7 @@ from openbb_charting.core.backend import create_backend, get_backend
 from openbb_cli.config.constants import AVAILABLE_FLAIRS, ENV_FILE_SETTINGS
 from openbb_cli.session import Session
 from openbb_core.app.model.charts.charting_settings import ChartingSettings
+from openbb_core.app.model.obbject import OBBject
 from pytz import all_timezones, timezone
 from rich.table import Table
 
@@ -974,3 +975,76 @@ def request(
         timeout=timeout,
         **kwargs,
     )
+
+
+def parse_unknown_args_to_dict(unknown_args: Optional[List[str]]) -> Dict[str, str]:
+    """Parse unknown arguments to a dictionary."""
+    unknown_args_dict = {}
+    if unknown_args:
+        for idx, arg in enumerate(unknown_args):
+            if arg.startswith("--"):
+                if idx + 1 < len(unknown_args):
+                    try:
+                        unknown_args_dict[arg.replace("--", "")] = (
+                            eval(  # noqa: S307, E501 pylint: disable=eval-used
+                                unknown_args[idx + 1]
+                            )
+                        )
+                    except Exception:
+                        unknown_args_dict[arg] = unknown_args[idx + 1]
+                else:
+                    session.console.print(
+                        f"Missing value for argument {arg}. Skipping this argument."
+                    )
+    return unknown_args_dict
+
+
+def handle_obbject_display(
+    obbject: OBBject,
+    chart: bool = False,
+    export: str = "",
+    sheet_name: str = "",
+    **kwargs,
+):
+    """Handle the display of an OBBject."""
+    df: pd.DataFrame = pd.DataFrame()
+    fig: Optional[OpenBBFigure] = None
+    if chart:
+        try:
+            if obbject.chart:
+                obbject.show(**kwargs)
+            else:
+                obbject.charting.to_chart(**kwargs)
+            if export:
+                fig = obbject.chart.fig
+                df = obbject.to_dataframe()
+        except Exception as e:
+            session.console.print(f"Failed to display chart: {e}")
+    else:
+        df = obbject.to_dataframe()
+        print_rich_table(
+            df=df,
+            show_index=True,
+            title=obbject.extra.get("command", ""),
+            export=bool(export),
+        )
+    if export and not df.empty:
+        if sheet_name and isinstance(sheet_name, list):
+            sheet_name = sheet_name[0]
+
+        func_name = (
+            obbject.extra.get("command", "")
+            .replace("/", "_")
+            .replace(" ", "_")
+            .replace("--", "_")
+        )
+        export_data(
+            export_type=",".join(export),
+            dir_path=os.path.dirname(os.path.abspath(__file__)),
+            func_name=func_name,
+            df=df,
+            sheet_name=sheet_name,
+            figure=fig,
+        )
+    elif export and df.empty:
+        session.console.print("[yellow]No data to export.[/yellow]")