Skip to content

Commit a95c555

Browse files
committed
Allow setting legend names when multiple series values are used in an Indicator kernc#382
1 parent 0a76e96 commit a95c555

File tree

3 files changed

+33
-6
lines changed

3 files changed

+33
-6
lines changed

backtesting/_plotting.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -522,10 +522,14 @@ def __eq__(self, other):
522522
colors = value._opts['color']
523523
colors = colors and cycle(_as_list(colors)) or (
524524
cycle([next(ohlc_colors)]) if is_overlay else colorgen())
525-
legend_label = LegendStr(value.name)
525+
legends = value._opts['legends']
526+
legends = legends and cycle(_as_list(legends))
527+
indicator_name = value.name
528+
legend_label = LegendStr(indicator_name)
526529
for j, arr in enumerate(value, 1):
527530
color = next(colors)
528-
source_name = f'{legend_label}_{i}_{j}'
531+
legend_label = next(legends) if legends is not None else legend_label
532+
source_name = f'{indicator_name}_{i}_{j}'
529533
if arr.dtype == bool:
530534
arr = arr.astype(int)
531535
source.add(arr, source_name)
@@ -563,9 +567,10 @@ def __eq__(self, other):
563567
line_color='#666666', line_dash='dashed',
564568
line_width=.5))
565569
if is_overlay:
566-
ohlc_tooltips.append((legend_label, NBSP.join(tooltips)))
570+
ohlc_tooltips.append((indicator_name, NBSP.join(tooltips)))
567571
else:
568-
set_tooltips(fig, [(legend_label, NBSP.join(tooltips))], vline=True, renderers=[r])
572+
set_tooltips(fig, [(indicator_name, NBSP.join(tooltips))],
573+
vline=True, renderers=[r])
569574
# If the sole indicator line on this figure,
570575
# have the legend only contain text without the glyph
571576
if len(value) == 1:

backtesting/backtesting.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ def _check_params(self, params):
7676
def I(self, # noqa: E741, E743
7777
func: Callable, *args,
7878
name=None, plot=True, overlay=None, color=None, scatter=False,
79-
**kwargs) -> np.ndarray:
79+
legends=None, **kwargs) -> np.ndarray:
8080
"""
8181
Declare indicator. An indicator is just an array of values,
8282
but one that is revealed gradually in
@@ -105,6 +105,10 @@ def I(self, # noqa: E741, E743
105105
If `scatter` is `True`, the plotted indicator marker will be a
106106
circle instead of a connected line segment (default).
107107
108+
`legends` can be list or array of string values to represent
109+
legends on your indicator chart. By default it's set to None,
110+
and `name` is used as legends.
111+
108112
Additional `*args` and `**kwargs` are passed to `func` and can
109113
be used for parameters.
110114
@@ -151,7 +155,7 @@ def init():
151155
overlay = ((x < 1.4) & (x > .6)).mean() > .6
152156

153157
value = _Indicator(value, name=name, plot=plot, overlay=overlay,
154-
color=color, scatter=scatter,
158+
color=color, scatter=scatter, legends=legends,
155159
# _Indicator.s Series accessor uses this:
156160
index=self.data.index)
157161
self._indicators.append(value)

backtesting/test/_test.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -771,6 +771,24 @@ def next(self):
771771
plot_drawdown=False, plot_equity=False, plot_pl=False, plot_volume=False,
772772
open_browser=False)
773773

774+
def test_indicator_legends(self):
775+
class S(Strategy):
776+
def init(self):
777+
self.I(lambda: (SMA(self.data.Close, 5), SMA(self.data.Close, 10)), overlay=False,
778+
name='Simple Moving Averages', scatter=False, legends=['SMA 5', 'SMA 10'])
779+
780+
def next(self):
781+
pass
782+
783+
bt = Backtest(GOOG, S)
784+
bt.run()
785+
with _tempfile() as f:
786+
bt.plot(filename=f,
787+
plot_drawdown=False, plot_equity=False, plot_pl=False, plot_volume=False,
788+
open_browser=True)
789+
# Give browser time to open before tempfile is removed
790+
time.sleep(1)
791+
774792

775793
class TestLib(TestCase):
776794
def test_barssince(self):

0 commit comments

Comments
 (0)