Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions tests/test_zipstream.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

import os
import tempfile
import time
import unittest
import zipstream
import zipfile
Expand Down Expand Up @@ -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')

Expand Down
12 changes: 8 additions & 4 deletions zipstream/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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:
Expand All @@ -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
Expand Down