def test_gzip_stream(self):
_, fname = self.create_file()
self.assertEqual(
un_gzip_stream(gzip_file(fname)).read(), b'hello world')
_, dirname = self.create_directory()
self.assertEqual(
un_gzip_stream(gzip_file(f"{dirname}/README.md")).read(),
b'hello world')
def test_gzip_stream(self):
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
self.addCleanup(lambda: os.remove(temp_file.name))
temp_file.write(b'contents')
name = temp_file.name
self.assertEqual(un_gzip_stream(gzip_file(name)).read(), b'contents')
def unpack(ext: str, source: IO[bytes], dest_path: str):
"""Unpack the archive |source| to |dest_path|.
Args:
ext (str): Extension of the archive.
source (IO[bytes]): File handle to the source.
dest_path ([type]): Destination path to unpack to.
"""
try:
if ext == '.tar.gz' or ext == '.tgz':
un_tar_directory(source, dest_path, 'gz')
elif ext == '.tar.bz2':
un_tar_directory(source, dest_path, 'bz2')
elif ext == '.bz2':
un_bz2_file(source, dest_path)
elif ext == '.gz':
with open(dest_path, 'wb') as f:
shutil.copyfileobj(un_gzip_stream(source), f)
elif ext == '.zip':
unzip_directory(source, dest_path)
else:
raise UsageError('Not an archive.')
except (tarfile.TarError, IOError) as e:
logging.error("Invalid archive upload: failed to unpack archive: %s",
e)
raise UsageError('Invalid archive upload: failed to unpack archive.')
def stream_file(self, target, gzipped):
"""
Returns a file-like object reading the given file. This file is gzipped
if gzipped is True.
"""
if self._is_available_locally(target):
file_path = self._get_target_path(target)
if gzipped:
return self.file_util.gzip_file(file_path)
else:
with self.file_util.OpenFile(file_path, gzipped=False) as f:
return f
else:
worker = self._bundle_model.get_bundle_worker(target.bundle_uuid)
response_socket_id = self._worker_model.allocate_socket(
worker['user_id'], worker['worker_id'])
try:
read_args = {'type': 'stream_file'}
self._send_read_message(worker, response_socket_id, target,
read_args)
fileobj = self._get_read_response_stream(response_socket_id)
if not gzipped:
fileobj = un_gzip_stream(fileobj)
return Deallocating(fileobj, self._worker_model,
response_socket_id)
except Exception:
self._worker_model.deallocate_socket(response_socket_id)
raise