add __enter__ and __exit__ to PdfPages

megies · megies · commit 5cb28403e5dd · 2013-09-12T15:22:27.000+02:00
diff --git a/lib/matplotlib/backends/backend_pdf.py b/lib/matplotlib/backends/backend_pdf.py
@@ -2258,6 +2258,12 @@ class PdfPages(object):
         # Once you are done, remember to close the object:
         pp.close()
 
+    Or using the with statement::
+
+        with PdfPages('foo.pdf') as pp:
+            fig.savefig(pp, format='pdf') # note the format argument!
+            pp.savefig(fig)
+
     (In reality PdfPages is a thin wrapper around PdfFile, in order to
     avoid confusion when using savefig and forgetting the format
     argument.)
@@ -2272,6 +2278,12 @@ def __init__(self, filename):
         """
         self._file = PdfFile(filename)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
     def close(self):
         """
         Finalize this object, making the underlying file a complete
