Merge pull request matplotlib#2440 from megies/pdfpages_pagecount

WeatherGod · WeatherGod · commit fc94ddbe4b26 · 2013-09-19T07:17:53.000-07:00
Pdfpages pagecount convenience getter method
diff --git a/lib/matplotlib/backends/backend_pdf.py b/lib/matplotlib/backends/backend_pdf.py
@@ -2310,6 +2310,12 @@ def savefig(self, figure=None, **kwargs):
             else:
                 figureManager.canvas.figure.savefig(self, format='pdf', **kwargs)
 
+    def get_pagecount(self):
+        """
+        Returns the current number of pages in the multipage pdf file.
+        """
+        return len(self._file.pageList)
+
 class FigureCanvasPdf(FigureCanvasBase):
     """
     The canvas the figure renders into.  Calls the draw and print fig
diff --git a/lib/matplotlib/tests/test_backend_pdf.py b/lib/matplotlib/tests/test_backend_pdf.py
@@ -37,3 +37,23 @@ def test_type42():
     ax = fig.add_subplot(111)
     ax.plot([1,2,3])
     fig.savefig(io.BytesIO())
+
+
+@cleanup
+def test_multipage_pagecount():
+    from matplotlib.backends.backend_pdf import PdfPages
+    from io import BytesIO
+    with PdfPages(BytesIO()) as pdf:
+        assert pdf.get_pagecount() == 0
+        fig = plt.figure()
+        ax = fig.add_subplot(111)
+        ax.plot([1, 2, 3])
+        fig.savefig(pdf, format="pdf")
+        assert pdf.get_pagecount() == 1
+        pdf.savefig()
+        assert pdf.get_pagecount() == 2
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
