diff --git a/gamestonk_terminal/common/quantitative_analysis/qa_view.py b/gamestonk_terminal/common/quantitative_analysis/qa_view.py index dd085951f6d..da83657dc74 100644 --- a/gamestonk_terminal/common/quantitative_analysis/qa_view.py +++ b/gamestonk_terminal/common/quantitative_analysis/qa_view.py @@ -27,6 +27,7 @@ from gamestonk_terminal.helper_funcs import ( export_data, plot_autoscale, rich_table_from_df, + LineAnnotateDrawer, ) register_matplotlib_converters() @@ -563,7 +564,11 @@ def display_raw( def display_line( - data: pd.Series, title: str = "", log_y: bool = True, export: str = "" + data: pd.Series, + title: str = "", + log_y: bool = True, + draw: bool = False, + export: str = "", ): """Display line plot of data @@ -575,6 +580,8 @@ def display_line( Title for plot log_y: bool Flag for showing y on log scale + draw: bool + Flag for drawing lines and annotating on the plot export: str Format to export data """ @@ -597,8 +604,13 @@ def display_line( if title: fig.suptitle(title) fig.tight_layout(pad=2) + if gtff.USE_ION: plt.ion() + + if draw: + LineAnnotateDrawer(ax).draw_lines_and_annotate() + plt.show() export_data( diff --git a/gamestonk_terminal/helper_funcs.py b/gamestonk_terminal/helper_funcs.py index f83585d5f89..04492f58e9d 100644 --- a/gamestonk_terminal/helper_funcs.py +++ b/gamestonk_terminal/helper_funcs.py @@ -918,6 +918,46 @@ def try_except(f): return inner +class LineAnnotateDrawer: + def __init__(self, ax: matplotlib.axes = None): + self.ax = ax + + def draw_lines_and_annotate(self): + + # ymin, _ = self.ax.get_ylim() + # xmin, _ = self.ax.get_xlim() + # self.ax.plot( + # [xmin, xmin], + # [ymin, ymin], + # lw=0, + # color="white", + # label="X - leave interactive mode\nClick twice for annotation", + # ) + # self.ax.legend(handlelength=0, handletextpad=0, fancybox=True, loc=2) + # self.ax.figure.canvas.draw() + + print("Click twice for annotation.\nClose window to keep using terminal.\n") + + while True: + xy = plt.ginput(2) + # Check whether the user has closed the window or not + if not plt.get_fignums(): + print("") + return + + if len(xy) == 2: + x = [p[0] for p in xy] + y = [p[1] for p in xy] + + if (x[0] == x[1]) and (y[0] == y[1]): + txt = input("Annotation: ") + self.ax.annotate(txt, (x[0], y[1]), ha="center", va="center") + else: + self.ax.plot(x, y) + + self.ax.figure.canvas.draw() + + def system_clear(): """Clear screen""" os.system("cls||clear") # nosec diff --git a/gamestonk_terminal/stocks/quantitative_analysis/qa_controller.py b/gamestonk_terminal/stocks/quantitative_analysis/qa_controller.py index b8129619a52..df0f42c989c 100644 --- a/gamestonk_terminal/stocks/quantitative_analysis/qa_controller.py +++ b/gamestonk_terminal/stocks/quantitative_analysis/qa_controller.py @@ -32,6 +32,8 @@ from gamestonk_terminal.stocks.quantitative_analysis.factors_view import capm_vi t_console = Console() +# pylint: disable=C0302 + class QaController: """Quantitative Analysis Controller class""" @@ -439,6 +441,14 @@ Other: action="store_true", default=False, ) + parser.add_argument( + "-d", + "--draw", + help="Draw lines and annotate on the plot", + dest="draw", + action="store_true", + default=False, + ) ns_parser = parse_known_args_and_warn( parser, other_args, export_allowed=EXPORT_ONLY_FIGURES_ALLOWED @@ -448,6 +458,7 @@ Other: self.stock[self.target], title=f"{self.ticker} {self.target}", log_y=ns_parser.log, + draw=ns_parser.draw, ) @try_except