Merge pull request matplotlib#2453 from megies/pdfpages_keep_empty

PdfPages: add option to delete empty file when closed

Author: mdboom

lib/matplotlib/backends/backend_pdf.py

@@ -2251,29 +2251,44 @@ class PdfPages(object):
     """
     A multi-page PDF file.
 
-    Use like this::
-
-        # Initialize:
-        with PdfPages('foo.pdf') as pdf:
-
-            # As many times as you like, create a figure fig and save it:
-            # When no figure is specified the current figure is saved
-            pdf.savefig(fig)
-            pdf.savefig()
-
-    (In reality PdfPages is a thin wrapper around PdfFile, in order to
-    avoid confusion when using savefig and forgetting the format
-    argument.)
+    Examples
+    --------
+
+    >>> import matplotlib.pyplot as plt
+    >>> # Initialize:
+    >>> with PdfPages('foo.pdf') as pdf:
+    ...     # As many times as you like, create a figure fig and save it:
+    ...     fig = plt.figure()
+    ...     pdf.savefig(fig)
+    ...     # When no figure is specified the current figure is saved
+    ...     pdf.savefig()
+
+    Notes
+    -----
+
+    In reality :class:`PdfPages` is a thin wrapper around :class:`PdfFile`, in
+    order to avoid confusion when using :func:`~matplotlib.pyplot.savefig` and
+    forgetting the format argument.
     """
-    __slots__ = ('_file',)
+    __slots__ = ('_file', 'keep_empty')
 
-    def __init__(self, filename):
+    def __init__(self, filename, keep_empty=True):
         """
-        Create a new PdfPages object that will be written to the file
-        named *filename*. The file is opened at once and any older
-        file with the same name is overwritten.
+        Create a new PdfPages object.
+
+        Parameters
+        ----------
+
+        filename: str
+            Plots using :meth:`PdfPages.savefig` will be written to a file at
+            this location. The file is opened at once and any older file with
+            the same name is overwritten.
+        keep_empty: bool, optional
+            If set to False, then empty pdf files will be deleted automatically
+            when closed.
         """
         self._file = PdfFile(filename)
+        self.keep_empty = keep_empty
 
     def __enter__(self):
         return self
@@ -2287,6 +2302,8 @@ def close(self):
         PDF file.
         """
         self._file.close()
+        if self.get_pagecount() == 0 and self.keep_empty is False:
+            os.remove(self._file.fh.name)
         self._file = None
 
     def infodict(self):
@@ -2299,10 +2316,19 @@ def infodict(self):
 
     def savefig(self, figure=None, **kwargs):
         """
-        Save the Figure instance *figure* to this file as a new page.
-        If *figure* is a number, the figure instance is looked up by
-        number, and if *figure* is None, the active figure is saved.
-        Any other keyword arguments are passed to Figure.savefig.
+        Saves a :class:`~matplotlib.figure.Figure` to this file as a new page.
+
+        Any other keyword arguments are passed to
+        :meth:`~matplotlib.figure.Figure.savefig`.
+
+        Parameters
+        ----------
+
+        figure: :class:`~matplotlib.figure.Figure` or int, optional
+            Specifies what figure is saved to file. If not specified, the
+            active figure is saved. If a :class:`~matplotlib.figure.Figure`
+            instance is provided, this figure is saved. If an int is specified,
+            the figure instance to save is looked up by number.
         """
         if isinstance(figure, Figure):
             figure.savefig(self, format='pdf', **kwargs)

lib/matplotlib/tests/test_backend_pdf.py

@@ -51,3 +51,39 @@ def test_multipage_pagecount():
         assert pdf.get_pagecount() == 1
         pdf.savefig()
         assert pdf.get_pagecount() == 2
+
+
+@cleanup
+def test_multipage_keep_empty():
+    from matplotlib.backends.backend_pdf import PdfPages
+    from tempfile import NamedTemporaryFile
+    ### test empty pdf files
+    # test that an empty pdf is left behind with keep_empty=True (default)
+    with NamedTemporaryFile(delete=False) as tmp:
+        with PdfPages(tmp) as pdf:
+            filename = pdf._file.fh.name
+        assert os.path.exists(filename)
+        os.remove(filename)
+    # test if an empty pdf is deleting itself afterwards with keep_empty=False
+    with NamedTemporaryFile(delete=False) as tmp:
+        with PdfPages(tmp, keep_empty=False) as pdf:
+            filename = pdf._file.fh.name
+        assert not os.path.exists(filename)
+    ### test pdf files with content, they should never be deleted
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.plot([1, 2, 3])
+    # test that a non-empty pdf is left behind with keep_empty=True (default)
+    with NamedTemporaryFile(delete=False) as tmp:
+        with PdfPages(tmp) as pdf:
+            filename = pdf._file.fh.name
+            pdf.savefig()
+        assert os.path.exists(filename)
+        os.remove(filename)
+    # test that a non-empty pdf is left behind with keep_empty=False
+    with NamedTemporaryFile(delete=False) as tmp:
+        with PdfPages(tmp, keep_empty=False) as pdf:
+            filename = pdf._file.fh.name
+            pdf.savefig()
+        assert os.path.exists(filename)
+        os.remove(filename)