Added a way to post-process subplots.

pyalgotrade/plotter.py

@@ -53,6 +53,13 @@ def includeDateTime(self, dateTime):
     return filter(lambda x: dateTimeFilter.includeDateTime(x), dateTimes)
 
 
+def _post_plot_fun(subPlot, mplSubplot):
+    # Legend
+    mplSubplot.legend(subPlot.getAllSeries().keys(), shadow=True, loc="best")
+    # Don't scale the Y axis
+    mplSubplot.yaxis.set_major_formatter(ticker.ScalarFormatter(useOffset=False))
+
+
 class Series(object):
     def __init__(self):
         self.__values = {}
@@ -215,6 +222,9 @@ def __getColor(self, series):
     def isEmpty(self):
         return len(self.__series) == 0
 
+    def getAllSeries(self):
+        return self.__series
+
     def addDataSeries(self, label, dataSeries, defaultClass=LineMarker):
         """Add a DataSeries to the subplot.
 
@@ -261,20 +271,14 @@ def getSeries(self, name, defaultClass=LineMarker):
     def getCustomMarksSeries(self, name):
         return self.getSeries(name, CustomMarker)
 
-    def customizeSubplot(self, mplSubplot):
-        # Don't scale the Y axis
-        mplSubplot.yaxis.set_major_formatter(ticker.ScalarFormatter(useOffset=False))
-
-    def plot(self, mplSubplot, dateTimes):
+    def plot(self, mplSubplot, dateTimes, postPlotFun=_post_plot_fun):
         for series in self.__series.values():
             color = None
             if series.needColor():
                 color = self.__getColor(series)
             series.plot(mplSubplot, dateTimes, color)
 
-        # Legend
-        mplSubplot.legend(self.__series.keys(), shadow=True, loc="best")
-        self.customizeSubplot(mplSubplot)
+        postPlotFun(self, mplSubplot)
 
 
 class InstrumentSubplot(Subplot):
@@ -397,7 +401,7 @@ def getPortfolioSubplot(self):
         """
         return self.__portfolioSubplot
 
-    def __buildFigureImpl(self, fromDateTime=None, toDateTime=None):
+    def __buildFigureImpl(self, fromDateTime=None, toDateTime=None, postPlotFun=_post_plot_fun):
         dateTimes = _filter_datetimes(self.__dateTimes, fromDateTime, toDateTime)
         dateTimes.sort()
 
@@ -414,7 +418,7 @@ def __buildFigureImpl(self, fromDateTime=None, toDateTime=None):
             axesSubplot = axes[i][0]
             if not subplot.isEmpty():
                 mplSubplots.append(axesSubplot)
-                subplot.plot(axesSubplot, dateTimes)
+                subplot.plot(axesSubplot, dateTimes, postPlotFun=postPlotFun)
                 axesSubplot.grid(True)
 
         return (fig, mplSubplots)
@@ -426,7 +430,7 @@ def buildFigure(self, fromDateTime=None, toDateTime=None):
         fig, _ = self.buildFigureAndSubplots(fromDateTime, toDateTime)
         return fig
 
-    def buildFigureAndSubplots(self, fromDateTime=None, toDateTime=None):
+    def buildFigureAndSubplots(self, fromDateTime=None, toDateTime=None, postPlotFun=_post_plot_fun):
         """Builds a matplotlib.figure.Figure with the subplots. Must be called after running the strategy.
 
         :param fromDateTime: An optional starting datetime.datetime. Everything before it won't get plotted.
@@ -435,11 +439,11 @@ def buildFigureAndSubplots(self, fromDateTime=None, toDateTime=None):
         :type toDateTime: datetime.datetime
         :rtype: A 2 element tuple with matplotlib.figure.Figure and subplots.
         """
-        fig, mplSubplots = self.__buildFigureImpl(fromDateTime, toDateTime)
+        fig, mplSubplots = self.__buildFigureImpl(fromDateTime, toDateTime, postPlotFun=postPlotFun)
         fig.autofmt_xdate()
         return fig, mplSubplots
 
-    def plot(self, fromDateTime=None, toDateTime=None):
+    def plot(self, fromDateTime=None, toDateTime=None, postPlotFun=_post_plot_fun):
         """Plots the strategy execution. Must be called after running the strategy.
 
         :param fromDateTime: An optional starting datetime.datetime. Everything before it won't get plotted.
@@ -448,6 +452,6 @@ def plot(self, fromDateTime=None, toDateTime=None):
         :type toDateTime: datetime.datetime
         """
 
-        fig, mplSubplots = self.__buildFigureImpl(fromDateTime, toDateTime)
+        fig, mplSubplots = self.__buildFigureImpl(fromDateTime, toDateTime, postPlotFun=postPlotFun)
         fig.autofmt_xdate()
         plt.show()