Skip to content

Commit

Permalink
Allow setting legend names when multiple series values are used in an…
Browse files Browse the repository at this point in the history
… Indicator kernc#382
  • Loading branch information
zlpatel committed Jun 16, 2021
1 parent 0a76e96 commit 31bc463
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 6 deletions.
13 changes: 9 additions & 4 deletions backtesting/_plotting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
8 changes: 6 additions & 2 deletions backtesting/backtesting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
18 changes: 18 additions & 0 deletions backtesting/test/_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 31bc463

Please sign in to comment.