New datetime parameter into write_iter.

Author: oszkarkmetti

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,9 +220,9 @@ 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):
@@ -233,7 +233,7 @@ def _iterable():
             yield data
         return self.write_iter(arcname, _iterable(), compress_type=compress_type, buffer_size=buffer_size)
 
-    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