Support writing iterable streams

This commit is contained in:
Chris AtLee 2017-08-23 11:50:50 -04:00
Родитель df58bf063c
Коммит 7642fc8909
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: 27EF48554428DB0C
1 изменённых файлов: 22 добавлений и 8 удалений

Просмотреть файл

@ -137,27 +137,40 @@ class MarWriter(object):
for f in files: for f in files:
self.add_file(os.path.join(root, f), compress) self.add_file(os.path.join(root, f), compress)
def add_fileobj(self, fileobj, path, compress): def add_fileobj(self, fileobj, path, compress, flags=None):
"""Add the contents of a file object to the MAR file. """Add the contents of a file object to the MAR file.
Args: Args:
fileobj (file-like object): open file object data will be read from fileobj (file-like object): open file object
path (str): name of this file in the MAR file path (str): name of this file in the MAR file
compress (str): One of 'xz', 'bz2', or None. Defaults to None. compress (str): One of 'xz', 'bz2', or None. Defaults to None.
flags (int): permission of this file in the MAR file. Defaults to the permissions of `path`
"""
f = file_iter(fileobj)
flags = flags or os.stat(path) & 0o777
return self.add_stream(f, path, compress, flags)
def add_stream(self, stream, path, compress, flags):
"""Add the contents of an iterable to the MAR file.
Args:
stream (iterable): yields blocks of data
path (str): name of this file in the MAR file
compress (str): One of 'xz', 'bz2', or None. Defaults to None.
flags (int): permission of this file in the MAR file
""" """
self.data_fileobj.seek(self.last_offset) self.data_fileobj.seek(self.last_offset)
f = file_iter(fileobj)
if compress == 'bz2': if compress == 'bz2':
f = bz2_compress_stream(f) stream = bz2_compress_stream(stream)
elif compress == 'xz': elif compress == 'xz':
f = xz_compress_stream(f) stream = xz_compress_stream(stream)
elif compress is None: elif compress is None:
pass pass
else: else:
raise ValueError('Unsupported compression type: {}'.format(compress)) raise ValueError('Unsupported compression type: {}'.format(compress))
size = write_to_file(f, self.data_fileobj) size = write_to_file(stream, self.data_fileobj)
# On Windows, convert \ to / # On Windows, convert \ to /
# very difficult to mock this out for coverage on linux # very difficult to mock this out for coverage on linux
@ -168,7 +181,7 @@ class MarWriter(object):
name=path, name=path,
offset=self.last_offset, offset=self.last_offset,
size=size, size=size,
flags=os.stat(path).st_mode & 0o777, flags=flags,
) )
self.entries.append(e) self.entries.append(e)
self.last_offset += e['size'] self.last_offset += e['size']
@ -185,7 +198,8 @@ class MarWriter(object):
self.fileobj.seek(self.last_offset) self.fileobj.seek(self.last_offset)
with open(path, 'rb') as f: with open(path, 'rb') as f:
self.add_fileobj(f, path, compress) flags = os.stat(path).st_mode & 0o777
self.add_fileobj(f, path, compress, flags)
def write_header(self): def write_header(self):
"""Write the MAR header to the file. """Write the MAR header to the file.