diff --git a/backtesting/_plotting.py b/backtesting/_plotting.py index bec6757a..368b14f6 100644 --- a/backtesting/_plotting.py +++ b/backtesting/_plotting.py @@ -522,10 +522,14 @@ def __eq__(self, other): colors = value._opts['color'] colors = colors and cycle(_as_list(colors)) or ( cycle([next(ohlc_colors)]) if is_overlay else colorgen()) - legend_label = LegendStr(value.name) + legends = value._opts['legends'] + legends = legends and cycle(_as_list(legends)) + indicator_name = value.name + legend_label = LegendStr(indicator_name) for j, arr in enumerate(value, 1): color = next(colors) - source_name = f'{legend_label}_{i}_{j}' + legend_label = next(legends) if legends is not None else legend_label + source_name = f'{indicator_name}_{i}_{j}' if arr.dtype == bool: arr = arr.astype(int) source.add(arr, source_name) @@ -563,9 +567,10 @@ def __eq__(self, other): line_color='#666666', line_dash='dashed', line_width=.5)) if is_overlay: - ohlc_tooltips.append((legend_label, NBSP.join(tooltips))) + ohlc_tooltips.append((indicator_name, NBSP.join(tooltips))) else: - set_tooltips(fig, [(legend_label, NBSP.join(tooltips))], vline=True, renderers=[r]) + set_tooltips(fig, [(indicator_name, NBSP.join(tooltips))], + vline=True, renderers=[r]) # If the sole indicator line on this figure, # have the legend only contain text without the glyph if len(value) == 1: diff --git a/backtesting/backtesting.py b/backtesting/backtesting.py index edb7be01..368c1402 100644 --- a/backtesting/backtesting.py +++ b/backtesting/backtesting.py @@ -76,7 +76,7 @@ def _check_params(self, params): def I(self, # noqa: E741, E743 func: Callable, *args, name=None, plot=True, overlay=None, color=None, scatter=False, - **kwargs) -> np.ndarray: + legends=None, **kwargs) -> np.ndarray: """ Declare indicator. An indicator is just an array of values, but one that is revealed gradually in @@ -105,6 +105,10 @@ def I(self, # noqa: E741, E743 If `scatter` is `True`, the plotted indicator marker will be a circle instead of a connected line segment (default). + `legends` can be list or array of string values to represent + legends on your indicator chart. By default it's set to None, + and `name` is used as legends. + Additional `*args` and `**kwargs` are passed to `func` and can be used for parameters. @@ -151,7 +155,7 @@ def init(): overlay = ((x < 1.4) & (x > .6)).mean() > .6 value = _Indicator(value, name=name, plot=plot, overlay=overlay, - color=color, scatter=scatter, + color=color, scatter=scatter, legends=legends, # _Indicator.s Series accessor uses this: index=self.data.index) self._indicators.append(value) diff --git a/backtesting/test/_test.py b/backtesting/test/_test.py index 85ecea6a..1ca822cf 100644 --- a/backtesting/test/_test.py +++ b/backtesting/test/_test.py @@ -771,6 +771,24 @@ def next(self): plot_drawdown=False, plot_equity=False, plot_pl=False, plot_volume=False, open_browser=False) + def test_indicator_legends(self): + class S(Strategy): + def init(self): + self.I(lambda: (SMA(self.data.Close, 5), SMA(self.data.Close, 10)), overlay=False, + name='Simple Moving Averages', scatter=False, legends=['SMA 5', 'SMA 10']) + + def next(self): + pass + + bt = Backtest(GOOG, S) + bt.run() + with _tempfile() as f: + bt.plot(filename=f, + plot_drawdown=False, plot_equity=False, plot_pl=False, plot_volume=False, + open_browser=True) + # Give browser time to open before tempfile is removed + time.sleep(1) + class TestLib(TestCase): def test_barssince(self):