plot.generate_candlestick_graph Added short equivelent, separating plotting scatter creation to a function

freqtrade/plot/plotting.py

@@ -1,6 +1,6 @@
 import logging
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import pandas as pd
 
@@ -385,6 +385,35 @@ def add_areas(fig, row: int, data: pd.DataFrame, indicators) -> make_subplots:
     return fig
 
 
+def create_scatter(
+    data,
+    column_name,
+    color,
+    direction
+) -> Optional[go.Scatter]:
+
+    if column_name in data.columns:
+        df_short = data[data[column_name] == 1]
+        if len(df_short) > 0:
+            shorts = go.Scatter(
+                x=df_short.date,
+                y=df_short.close,
+                mode='markers',
+                name=column_name,
+                marker=dict(
+                    symbol=f"triangle-{direction}-dot",
+                    size=9,
+                    line=dict(width=1),
+                    color=color,
+                )
+            )
+            return shorts
+        else:
+            logger.warning(f"No {column_name}-signals found.")
+
+    return None
+
+
 def generate_candlestick_graph(pair: str, data: pd.DataFrame, trades: pd.DataFrame = None, *,
                                indicators1: List[str] = [],
                                indicators2: List[str] = [],
@@ -431,44 +460,15 @@ def generate_candlestick_graph(pair: str, data: pd.DataFrame, trades: pd.DataFra
     )
     fig.add_trace(candles, 1, 1)
 
-    # TODO-lev: Needs short equivalent
-    if 'enter_long' in data.columns:
-        df_buy = data[data['enter_long'] == 1]
-        if len(df_buy) > 0:
-            buys = go.Scatter(
-                x=df_buy.date,
-                y=df_buy.close,
-                mode='markers',
-                name='buy',
-                marker=dict(
-                    symbol='triangle-up-dot',
-                    size=9,
-                    line=dict(width=1),
-                    color='green',
-                )
-            )
-            fig.add_trace(buys, 1, 1)
-        else:
-            logger.warning("No buy-signals found.")
+    longs = create_scatter(data, 'enter_long', 'green', 'up')
+    exit_longs = create_scatter(data, 'exit_long', 'red', 'down')
+    shorts = create_scatter(data, 'enter_short', 'blue', 'down')
+    exit_shorts = create_scatter(data, 'exit_short', 'violet', 'up')
+
+    for scatter in [longs, exit_longs, shorts, exit_shorts]:
+        if scatter:
+            fig.add_trace(scatter, 1, 1)
 
-    if 'exit_long' in data.columns:
-        df_sell = data[data['exit_long'] == 1]
-        if len(df_sell) > 0:
-            sells = go.Scatter(
-                x=df_sell.date,
-                y=df_sell.close,
-                mode='markers',
-                name='sell',
-                marker=dict(
-                    symbol='triangle-down-dot',
-                    size=9,
-                    line=dict(width=1),
-                    color='red',
-                )
-            )
-            fig.add_trace(sells, 1, 1)
-        else:
-            logger.warning("No sell-signals found.")
     # Add Bollinger Bands
     fig = plot_area(fig, 1, data, 'bb_lowerband', 'bb_upperband',
                     label="Bollinger Band")
@@ -537,7 +537,7 @@ def generate_profit_graph(pairs: str, data: Dict[str, pd.DataFrame],
                             "Profit per pair",
                             "Parallelism",
                             "Underwater",
-                            ])
+                        ])
     fig['layout'].update(title="Freqtrade Profit plot")
     fig['layout']['yaxis1'].update(title='Price')
     fig['layout']['yaxis2'].update(title=f'Profit {stake_currency}')