Use ContextualZipFile for context manager support

1 files changed, 22 insertions, 5 deletions

diff --git a/pkg_resources.py b/pkg_resources.py
index 63a31fb8..ba03f628 100644
--- a/pkg_resources.py
+++ b/pkg_resources.py
@@ -1551,18 +1551,35 @@ def build_zipmanifest(path):
       * [7] - zipinfo.CRC
     """
     zipinfo = dict()
-    zfile = zipfile.ZipFile(path)
-    #Got ZipFile has not __exit__ on python 3.1
-    try:
+    with ContextualZipFile.compat(path) as zfile:
         for zitem in zfile.namelist():
             zpath = zitem.replace('/', os.sep)
             zipinfo[zpath] = zfile.getinfo(zitem)
             assert zipinfo[zpath] is not None
-    finally:
-        zfile.close()
     return zipinfo
 
 
+class ContextualZipFile(zipfile.ZipFile):
+    """
+    Supplement ZipFile class to support context manager for Python 2.6
+    """
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
+    @classmethod
+    def compat(cls, *args, **kwargs):
+        """
+        Construct a ZipFile or ContextualZipFile as appropriate
+        """
+        zf_has_exit = hasattr(zipfile.ZipFile, '__exit__')
+        class_ = zipfile.ZipFile if zf_has_exit else cls
+        return class_(*args, **kwargs)
+
+
 class ZipProvider(EggProvider):
     """Resource support for zips and eggs"""