Merge pull request #12 from allanlei/bug-7

Bug 7

example.py

@@ -8,6 +8,9 @@
     z.write('LICENSE')
     z.write('LICENSE', arcname='stuff/LICENSE')
 
+    with open('tests/sample.rtf', 'rb') as fp:
+        z.writestr('sample.rtf', fp.read())
+
     for root, directories, files in os.walk('zipstream'):
         for filename in files:
             path = os.path.join(root, filename)

tests/sample.rtf

@@ -0,0 +1,34 @@
+{\rtf1\ansi\deff3\adeflang1025
+{\fonttbl{\f0\froman\fprq2\fcharset0 Times New Roman;}{\f1\froman\fprq2\fcharset2 Symbol;}{\f2\fswiss\fprq2\fcharset0 Arial;}{\f3\froman\fprq2\fcharset0 Liberation Serif{\*\falt Times New Roman};}{\f4\fswiss\fprq2\fcharset0 Liberation Sans{\*\falt Arial};}{\f5\fmodern\fprq1\fcharset0 Liberation Mono{\*\falt Courier New};}{\f6\fnil\fprq2\fcharset0 Droid Sans Fallback;}{\f7\fmodern\fprq1\fcharset0 Courier New;}{\f8\fnil\fprq2\fcharset0 FreeSans;}{\f9\fswiss\fprq0\fcharset128 FreeSans;}}
+{\colortbl;\red0\green0\blue0;\red0\green0\blue255;\red0\green255\blue255;\red0\green255\blue0;\red255\green0\blue255;\red255\green0\blue0;\red255\green255\blue0;\red255\green255\blue255;\red0\green0\blue128;\red0\green128\blue128;\red0\green128\blue0;\red128\green0\blue128;\red128\green0\blue0;\red128\green128\blue0;\red128\green128\blue128;\red192\green192\blue192;}
+{\stylesheet{\s0\snext0\nowidctlpar\hyphpar0\cf0\dbch\af6\langfe2052\dbch\af8\afs24\alang1081\loch\f3\fs24\lang4105 Normal;}
+{\s1\sbasedon15\snext16\ilvl0\outlinelevel0\sb240\sa120\keepn\b\dbch\af6\dbch\af8\afs36\ab\loch\f4\fs36 Heading 1;}
+{\s2\sbasedon15\snext16\ilvl1\outlinelevel1\sb200\sa120\keepn\b\dbch\af6\dbch\af8\afs32\ab\loch\f4\fs32 Heading 2;}
+{\s3\sbasedon15\snext16\ilvl2\outlinelevel2\sb140\sa120\keepn\b\dbch\af6\dbch\af8\afs28\ab\loch\f4\fs28 Heading 3;}
+{\s15\sbasedon0\snext16\sb240\sa120\keepn\dbch\af6\dbch\af8\afs28\loch\f4\fs28 Heading;}
+{\s16\sbasedon0\snext16\sl288\slmult1\sb0\sa140 Text Body;}
+{\s17\sbasedon16\snext17\sl288\slmult1\sb0\sa140\dbch\af9 List;}
+{\s18\sbasedon0\snext18\sb120\sa120\noline\i\dbch\af9\afs24\ai\fs24 Caption;}
+{\s19\sbasedon0\snext19\noline\dbch\af9 Index;}
+{\s20\sbasedon0\snext20\sb0\sa0\dbch\af7\dbch\af5\afs20\loch\f5\fs20 Preformatted Text;}
+{\s21\sbasedon0\snext21\li567\ri567\lin567\rin567\fi0\sb0\sa283 Quotations;}
+{\s22\sbasedon15\snext16\qc\sb240\sa120\keepn\b\dbch\af6\dbch\af8\afs56\ab\loch\f4\fs56 Title;}
+{\s23\sbasedon15\snext16\qc\sb60\sa120\keepn\dbch\af6\dbch\af8\afs36\loch\f4\fs36 Subtitle;}
+}{\*\listtable{\list\listtemplateid1
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-432\li432}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-576\li576}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-720\li720}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-864\li864}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-1008\li1008}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-1152\li1152}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-1296\li1296}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-1440\li1440}
+{\listlevel\levelnfc255\leveljc0\levelstartat1\levelfollow2{\leveltext \'00;}{\levelnumbers;}\fi-1584\li1584}\listid1}
+}{\listoverridetable{\listoverride\listid1\listoverridecount0\ls1}}{\*\generator LibreOffice/5.0.2.2$Linux_X86_64 LibreOffice_project/00m0$Build-2}{\info{\creatim\yr0\mo0\dy0\hr0\min0}{\revtim\yr2015\mo9\dy29\hr20\min12}{\printim\yr0\mo0\dy0\hr0\min0}}\deftab709
+
+{\*\pgdsctbl
+{\pgdsc0\pgdscuse451\pgwsxn12240\pghsxn15840\marglsxn1134\margrsxn1134\margtsxn1134\margbsxn1134\pgdscnxt0 Default Style;}}
+\formshade\paperh15840\paperw12240\margl1134\margr1134\margt1134\margb1134\sectd\sbknone\sectunlocked1\pgndec\pgwsxn12240\pghsxn15840\marglsxn1134\margrsxn1134\margtsxn1134\margbsxn1134\ftnbj\ftnstart1\ftnrstcont\ftnnar\aenddoc\aftnrstcont\aftnstart1\aftnnrlc
+{\*\ftnsep\chftnsep}\pgndec\pard\plain \s20\sb0\sa0\dbch\af7\dbch\af5\afs20\loch\f5\fs20{\rtlch \ltrch\loch
+Hello World}
+\par }

tests/test_zipstream.py

@@ -8,6 +8,9 @@
 import zipfile
 
 
+SAMPLE_FILE_RTF = 'tests/sample.rtf'
+
+
 class ZipInfoTestCase(unittest.TestCase):
     pass
 
@@ -73,6 +76,21 @@ def string_generator():
 
         os.remove(f.name)
 
+    def test_writestr(self):
+        z = zipstream.ZipFile(mode='w')
+
+        with open(SAMPLE_FILE_RTF, 'rb') as fp:
+            z.writestr('sample.rtf', fp.read())
+
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        for chunk in z:
+            f.write(chunk)
+        f.close()
+
+        z2 = zipfile.ZipFile(f.name, 'r')
+        self.assertFalse(z2.testzip())
+
+        os.remove(f.name)
 
     def test_write_iterable_no_archive(self):
         z = zipstream.ZipFile(mode='w')

zipstream/__init__.py

@@ -220,6 +220,14 @@ def write_iter(self, arcname, iterable, compress_type=None):
         kwargs = {'arcname': arcname, 'iterable': iterable, 'compress_type': compress_type}
         self.paths_to_write.append(kwargs)
 
+    def writestr(self, arcname, data, compress_type=None):
+        """
+        Writes a str into ZipFile by wrapping data as a generator
+        """
+        def _iterable():
+            yield data
+        return self.write_iter(arcname, _iterable(), compress_type=compress_type)
+
     def __write(self, filename=None, iterable=None, arcname=None, compress_type=None):
         """Put the bytes from filename into the archive under the name
         `arcname`."""