diff --git a/tests/test_zipstream.py b/tests/test_zipstream.py index 6bda736..41bb368 100644 --- a/tests/test_zipstream.py +++ b/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') diff --git a/zipstream/__init__.py b/zipstream/__init__.py index a086e3a..bc03bb3 100644 --- a/zipstream/__init__.py +++ b/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