From fe56a7896337ee3ca8284b8b1a6bc9b1e0520f73 Mon Sep 17 00:00:00 2001
From: Bob Lannon <blannon@gmail.com>
Date: Mon, 15 Apr 2013 17:09:56 -0400
Subject: [PATCH] added ZippedFiles object

---
 saucebrush/utils.py | 40 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/saucebrush/utils.py b/saucebrush/utils.py
index 4f95e5d..f94e15d 100644
--- a/saucebrush/utils.py
+++ b/saucebrush/utils.py
@@ -124,4 +124,42 @@ class RemoteFile(object):
         resp = urlopen(self._url)
         for line in resp:
             yield line.rstrip()
-        resp.close()
\ No newline at end of file
+        resp.close()
+
+class ZippedFiles(object):
+    """ unpack a zipped collection of files on init.
+
+        Takes a string with file location or zipfile.ZipFile object
+
+        Best to wrap this in a Files() object, if the goal is to have a
+        linereader, as this only returns filelike objects.
+
+        if using a ZipFile object, make sure to set mode to 'a' or 'w' in order
+        to use the add() function. 
+    """
+    def __init__(self, zippedfile):
+        import zipfile
+        if type(zippedfile) == str:
+            self._zipfile = zipfile.ZipFile(zippedfile,'a')
+        else:
+            self._zipfile = zippedfile
+        self.paths = self._zipfile.namelist()
+        self.file_open_callback = None
+
+    def __iter__(self):
+        return self.filereader()
+
+    def add(self, path, dirname=None, arcname=None):
+        path_base = os.path.basename(path)
+        if dirname:
+            arcname = os.path.join(dirname,path_base)
+        if not arcname:
+            arcname = path_base
+        self._zipfile.write(path,arcname)
+        self.paths.append(path)
+
+    def filereader(self):
+        for path in iter(self.paths):
+            if self.file_open_callback:
+                self.file_open_callback(path)
+            yield self._zipfile.open(path)