Refactor savefig functionality to save plots of multiple strategies and multiple figures (numfigs > 1)

mementum · mementum · commit 8bb3612475d1 · 2017-05-03T00:44:10.000+02:00
diff --git a/backtrader/cerebro.py b/backtrader/cerebro.py
@@ -865,7 +865,8 @@ def getbroker(self):
     broker = property(getbroker, setbroker)
 
     def plot(self, plotter=None, numfigs=1, iplot=True, start=None, end=None,
-             savefig=False, figfilename='backtrader-plot-%i.png',
+             savefig=False, figfilename='backtrader-plot-{j}-{i}.png',
+             width=16, height=9, dpi=300, tight=True,
              **kwargs):
         '''
         Plots the strategies inside cerebro
@@ -886,6 +887,21 @@ def plot(self, plotter=None, numfigs=1, iplot=True, start=None, end=None,
         ``end``: An index to the datetime line array of the strategy or a
         ``datetime.date``, ``datetime.datetime`` instance indicating the end
         of the plot
+
+        ``savefig``: set to ``True`` to save to a file rather than plot
+
+        ``figfilename``: name of the file. Use ``{j}`` in the name for the
+        strategy index to which the figure corresponds and use ``{i}`` to
+        insert figure number if multiple figures are being used per strategy
+        plot
+
+        ``width``: in inches of the saved figure
+
+        ``height``: in inches of the saved figure
+
+        ``dpi``: quality in dots per inches of the saved figure
+
+        ``tight``: only save actual content and not the frame of the figure
         '''
         if self._exactbars > 0:
             return
@@ -914,7 +930,12 @@ def plot(self, plotter=None, numfigs=1, iplot=True, start=None, end=None,
                 figs.append(rfig)
 
             if savefig:
-                plotter.savefig(figfilename % len(figs))
+                for j, stratfigs in enumerate(figs):
+                    for i, fig in enumerate(stratfigs):
+                        plotter.savefig(fig,
+                                        filename=figfilename.format(j=j, i=i),
+                                        width=width, height=height, dpi=dpi,
+                                        tight=tight)
             else:
                 plotter.show()
         return figs
diff --git a/backtrader/plot/plot.py b/backtrader/plot/plot.py
@@ -757,10 +757,10 @@ def plotdata(self, data, indicators):
     def show(self):
         self.mpyplot.show()
 
-    def savefig(self, filename='backtrader-plot.png'):
-        fig = self.mpyplot.gcf()
-        fig.set_size_inches(16, 9)
-        self.mpyplot.savefig(filename, dpi=300, bbox_inches='tight')
+    def savefig(self, fig, filename, width=16, height=9, dpi=300, tight=True):
+        fig.set_size_inches(width, height)
+        bbox_inches = 'tight' * tight or None
+        fig.savefig(filename, dpi=dpi, bbox_inches=bbox_inches)
 
     def sortdataindicators(self, strategy):
         # These lists/dictionaries hold the subplots that go above each data
