Merge pull request #8 from okmetti/add_datetime_to_write_iter

CHANGELOG.md

@@ -2,11 +2,14 @@
 
 This file details the changes that were made after forking v1.1.4 from https://github.com/allanlei/python-zipstream.
 
-## v1.1.5 (2019-03-18)
-* Support Zip64 when compressing iterables and strings (https://github.com/allanlei/python-zipstream/pull/25)
+## Unreleased
+* New datetime parameter in write_iter (https://github.com/arjan-s/python-zipstream/pull/8)
+
+## v1.1.7 (2019-10-22)
+* Stream data in the order it was received (https://github.com/arjan-s/python-zipstream/pull/4)
 
 ## v1.1.6 (2019-06-06)
 * Add partial flushing of ZipStreams (https://github.com/arjan-s/python-zipstream/pull/1)
 
-## v1.1.7 (2019-10-22)
-* Stream data in the order it was received (https://github.com/arjan-s/python-zipstream/pull/4)
+## v1.1.5 (2019-03-18)
+* Support Zip64 when compressing iterables and strings (https://github.com/allanlei/python-zipstream/pull/25)

tests/test_zipstream.py

@@ -3,6 +3,7 @@
 
 import os
 import tempfile
+import time
 import unittest
 import zipstream
 import zipfile
@@ -76,6 +77,28 @@ def string_generator():
 
         os.remove(f.name)
 
+    def test_write_iterable_with_date_time(self):
+        file_name_in_zip = "data_datetime"
+        file_date_time_in_zip = time.strptime("2011-04-19 22:30:21", "%Y-%m-%d %H:%M:%S")
+
+        z = zipstream.ZipFile(mode='w')
+        def string_generator():
+            for _ in range(10):
+                yield b'zipstream\x01\n'
+        z.write_iter(iterable=string_generator(), arcname=file_name_in_zip, date_time=file_date_time_in_zip)
+
+        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())
+
+        self.assertEqual(file_date_time_in_zip[0:5], z2.getinfo(file_name_in_zip).date_time[0:5])
+
+        os.remove(f.name)
+
     def test_writestr(self):
         z = zipstream.ZipFile(mode='w')
 

zipstream/__init__.py

@@ -220,20 +220,20 @@ def write(self, filename, arcname=None, compress_type=None):
         kwargs = {'filename': filename, 'arcname': arcname, 'compress_type': compress_type}
         self.paths_to_write.append(kwargs)
 
-    def write_iter(self, arcname, iterable, compress_type=None, buffer_size=None):
+    def write_iter(self, arcname, iterable, compress_type=None, buffer_size=None, date_time=None):
         """Write the bytes iterable `iterable` to the archive under the name `arcname`."""
-        kwargs = {'arcname': arcname, 'iterable': iterable, 'compress_type': compress_type, 'buffer_size': buffer_size}
+        kwargs = {'arcname': arcname, 'iterable': iterable, 'compress_type': compress_type, 'buffer_size': buffer_size, 'date_time': date_time}
         self.paths_to_write.append(kwargs)
 
-    def writestr(self, arcname, data, compress_type=None, buffer_size=None):
+    def writestr(self, arcname, data, compress_type=None, buffer_size=None, date_time=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, buffer_size=buffer_size)
+        return self.write_iter(arcname, _iterable(), compress_type=compress_type, buffer_size=buffer_size, date_time=date_time)
 
-    def __write(self, filename=None, iterable=None, arcname=None, compress_type=None, buffer_size=None):
+    def __write(self, filename=None, iterable=None, arcname=None, compress_type=None, buffer_size=None, date_time=None):
         """Put the bytes from filename into the archive under the name
         `arcname`."""
         if not self.fp:
@@ -248,7 +248,11 @@ def __write(self, filename=None, iterable=None, arcname=None, compress_type=None
             mtime = time.localtime(st.st_mtime)
             date_time = mtime[0:6]
         else:
-            st, isdir, date_time = None, False, time.localtime()[0:6]
+            st, isdir = None, False
+            if date_time is not None and isinstance(date_time, time.struct_time):
+                date_time = date_time[0:6]
+            if date_time is None:
+                date_time = time.localtime()[0:6]
         # Create ZipInfo instance to store file information
         if arcname is None:
             arcname = filename