Fix bug whereby iter_zip returned closed ZipExtFile instances

featureflow/bytestream.py

@@ -79,15 +79,22 @@ def iter_zip(fn):
         for info in zf.filelist:
             if not info.file_size:
                 continue
-            with zf.open(info.filename) as f:
-                yield ZipWrapper(f, info)
+            f = zf.open(info.filename)
+            yield ZipWrapper(f, info)
 
 
 class ZipWrapper(object):
     def __init__(self, zipfile, zipinfo):
         self.zipinfo = zipinfo
         self.zipfile = zipfile
 
+    def __enter__(self):
+        self.zipfile.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.zipfile.__exit__(exc_type, exc_val, exc_tb)
+
     @property
     def file_size(self):
         return self.zipinfo.file_size

featureflow/test_bytestream.py

@@ -1,4 +1,4 @@
-from .bytestream import BytesWithTotalLength, ByteStream, ZipWrapper
+from .bytestream import BytesWithTotalLength, ByteStream, ZipWrapper, iter_zip
 import unittest2
 import sys
 import tempfile
@@ -136,3 +136,17 @@ def test_right_increment(self):
         x = b'blah'
         x += BytesWithTotalLength(b'fake', 100)
         self.assertEqual(b'blahfake', x)
+
+
+class IterZipTests(unittest2.TestCase):
+    def test_iter_zip_yields_open_zip_files(self):
+        bio = BytesIO()
+        filename = 'test.dat'
+        with zipfile.ZipFile(bio, mode='w') as zf:
+            zf.writestr(filename, b'content')
+        bio.seek(0)
+
+        with list(iter_zip(bio))[0] as z:
+            self.assertFalse(
+                z.zipfile.closed, 'zipfile should be open, but was closed')
+            self.assertEqual(z.zipfile.read(), b'content')