Refactor export functionality concept (#772)

* Refactor export functionality

* improve docs around -h in parse args

* Add export functionality to ba controller

Co-authored-by: Meghan Hone <meghanhone@meghans-mbp.home>

2 files changed, 58 insertions, 16 deletions

diff --git a/gamestonk_terminal/helper_funcs.py b/gamestonk_terminal/helper_funcs.py
index 127d3b58c6d..d0fc6c31c6e 100644
--- a/gamestonk_terminal/helper_funcs.py
+++ b/gamestonk_terminal/helper_funcs.py
@@ -1,7 +1,6 @@
 """Helper functions"""
 __docformat__ = "numpy"
 import argparse
-from argparse import ArgumentError
 from typing import List
 from datetime import datetime, timedelta, time as Time
 import os
@@ -26,6 +25,11 @@ register_matplotlib_converters()
 if cfgPlot.BACKEND is not None:
     matplotlib.use(cfgPlot.BACKEND)
 
+NO_EXPORT = 0
+EXPORT_ONLY_RAW_DATA_ALLOWED = 1
+EXPORT_ONLY_FIGURES_ALLOWED = 2
+EXPORT_BOTH_RAW_DATA_AND_FIGURES = 3
+
 
 def check_int_range(mini: int, maxi: int):
     """Checks if argparse argument is an int between 2 values.
@@ -482,7 +486,11 @@ def patch_pandas_text_adjustment():
     pandas.io.formats.format.TextAdjustment.adjoin = text_adjustment_adjoin
 
 
-def parse_known_args_and_warn(parser: argparse.ArgumentParser, other_args: List[str]):
+def parse_known_args_and_warn(
+    parser: argparse.ArgumentParser,
+    other_args: List[str],
+    export_allowed: int = NO_EXPORT,
+):
     """Parses list of arguments into the supplied parser
 
     Parameters
@@ -491,6 +499,9 @@ def parse_known_args_and_warn(parser: argparse.ArgumentParser, other_args: List[
         Parser with predefined arguments
     other_args: List[str]
         List of arguments to parse
+    export_allowed: int
+        Choose from NO_EXPORT, EXPORT_ONLY_RAW_DATA_ALLOWED,
+        EXPORT_ONLY_FIGURES_ALLOWED and EXPORT_BOTH_RAW_DATA_AND_FIGURES
 
     Returns
     -------
@@ -500,17 +511,28 @@ def parse_known_args_and_warn(parser: argparse.ArgumentParser, other_args: List[
     parser.add_argument(
         "-h", "--help", action="store_true", help="show this help message"
     )
-    try:
+    if export_allowed > NO_EXPORT:
+        choices_export = []
+        help_export = "Export "
+
+        if export_allowed == EXPORT_ONLY_RAW_DATA_ALLOWED:
+            choices_export += ["csv", "json", "xlsx"]
+            help_export += "raw data into csv, json, xlsx "
+        if export_allowed > EXPORT_ONLY_RAW_DATA_ALLOWED:
+            choices_export += ["png", "jpg", "pdf", "svg"]
+            if export_allowed == EXPORT_BOTH_RAW_DATA_AND_FIGURES:
+                help_export += "or "
+            help_export += "figure into png, jpg, pdf, svg "
+
         parser.add_argument(
             "--export",
-            choices=["png", "jpg", "pdf", "svg", "csv", "json", "xlsx"],
+            choices=choices_export,
             default="",
             type=str,
             dest="export",
-            help="Export plot to png,jpg,pdf,svg file or export dataframe to csv,json,xlsx",
+            help=help_export,
         )
-    except ArgumentError:
-        pass
+
     if gtff.USE_CLEAR_AFTER_CMD:
         os.system("cls||clear")
 
diff --git a/gamestonk_terminal/stocks/behavioural_analysis/ba_controller.py b/gamestonk_terminal/stocks/behavioural_analysis/ba_controller.py
index 858b16a9dd6..c59219de78a 100644
--- a/gamestonk_terminal/stocks/behavioural_analysis/ba_controller.py
+++ b/gamestonk_terminal/stocks/behavioural_analysis/ba_controller.py
@@ -12,6 +12,8 @@ from colorama import Style
 import pandas as pd
 from gamestonk_terminal import feature_flags as gtff
 from gamestonk_terminal.helper_funcs import (
+    EXPORT_BOTH_RAW_DATA_AND_FIGURES,
+    EXPORT_ONLY_RAW_DATA_ALLOWED,
     get_flair,
     parse_known_args_and_warn,
     check_int_range,
@@ -509,7 +511,9 @@ SentimentInvestor:
             help="limit messages shown.",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_ONLY_RAW_DATA_ALLOWED
+            )
             if not ns_parser:
                 return
 
@@ -538,7 +542,9 @@ SentimentInvestor:
             help="starting date (format YYYY-MM-DD) from when we are interested in stock's mentions.",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return
             if self._check_ticker():
@@ -567,7 +573,9 @@ SentimentInvestor:
             help="number of regions to plot that show highest interest.",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return
             if self._check_ticker():
@@ -595,7 +603,9 @@ SentimentInvestor:
             help="number of top related queries to print.",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return
             if self._check_ticker():
@@ -623,7 +633,9 @@ SentimentInvestor:
             help="number of top rising related queries to print.",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return
             if self._check_ticker():
@@ -656,7 +668,9 @@ SentimentInvestor:
         )
 
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return
 
@@ -697,7 +711,9 @@ SentimentInvestor:
             help="number of days in the past to extract tweets.",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return
             if self._check_ticker():
@@ -732,7 +748,9 @@ SentimentInvestor:
             help="Export dataframe data to csv,json,xlsx file",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return
             if self._check_ticker():
@@ -763,7 +781,9 @@ SentimentInvestor:
             help="Export dataframe data to csv,json,xlsx file",
         )
         try:
-            ns_parser = parse_known_args_and_warn(parser, other_args)
+            ns_parser = parse_known_args_and_warn(
+                parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            )
             if not ns_parser:
                 return