From 6ce29fa7a89cd266ca9dd3e21ed40f9100eb70f6 Mon Sep 17 00:00:00 2001
From: Georg Brandl <georg@python.org>
Date: Sat, 30 Oct 2010 14:33:28 +0000
Subject: [PATCH] Fix test_mailbox by supporting context manager protocol for
 get_file() returns.

---
 Doc/library/mailbox.rst | 10 +++++++---
 Lib/mailbox.py          | 14 ++++++++++++++
 Misc/NEWS               |  3 +++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/Doc/library/mailbox.rst b/Doc/library/mailbox.rst
index b041d94953c..da01791f70e 100644
--- a/Doc/library/mailbox.rst
+++ b/Doc/library/mailbox.rst
@@ -180,15 +180,19 @@ Maildir, mbox, MH, Babyl, and MMDF.
    .. method:: get_file(key)
 
       Return a file-like representation of the message corresponding to *key*,
-      or raise a :exc:`KeyError` exception if no such message exists. The
-      file-like object behaves as if open in binary mode. This file should be
+      or raise a :exc:`KeyError` exception if no such message exists.  The
+      file-like object behaves as if open in binary mode.  This file should be
       closed once it is no longer needed.
 
+      .. versionadded:: 3.2
+         The file-like object supports the context manager protocol, so that
+         you can use a :keyword:`with` statement to automatically close it.
+
       .. note::
 
          Unlike other representations of messages, file-like representations are
          not necessarily independent of the :class:`Mailbox` instance that
-         created them or of the underlying mailbox. More specific documentation
+         created them or of the underlying mailbox.  More specific documentation
          is provided by each subclass.
 
 
diff --git a/Lib/mailbox.py b/Lib/mailbox.py
index e6f17353362..520463a8693 100644
--- a/Lib/mailbox.py
+++ b/Lib/mailbox.py
@@ -1827,6 +1827,8 @@ class _ProxyFile:
 
     def close(self):
         """Close the file."""
+        if hasattr(self._file, 'close'):
+            self._file.close()
         del self._file
 
     def _read(self, size, read_method):
@@ -1838,6 +1840,13 @@ class _ProxyFile:
         self._pos = self._file.tell()
         return result
 
+    def __enter__(self):
+        """Context manager protocol support."""
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+
 
 class _PartialFile(_ProxyFile):
     """A read-only wrapper of part of a file."""
@@ -1871,6 +1880,11 @@ class _PartialFile(_ProxyFile):
             size = remaining
         return _ProxyFile._read(self, size, read_method)
 
+    def close(self):
+        # do *not* close the underlying file object for partial files,
+        # since it's global to the mailbox object
+        del self._file
+
 
 def _lock_file(f, dotlock=True):
     """Lock file f using lockf and dot locking."""
diff --git a/Misc/NEWS b/Misc/NEWS
index afbf04c04b3..827640a5d48 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -57,6 +57,9 @@ Core and Builtins
 Library
 -------
 
+- Support context manager protocol for file-like objects returned by
+  mailbox ``get_file()`` methods.
+
 - Issue #10246: uu.encode didn't close file objects explicitly when filenames
   were given to it.  Patch by Brian Brazil.