Fix test_mailbox by supporting context manager protocol for get_file() returns.

1 files changed, 14 insertions, 0 deletions

diff --git a/Lib/mailbox.py b/Lib/mailbox.py
index e6f1735..520463a 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."""