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 test_gzip_stream(self):
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
self.addCleanup(lambda: os.remove(temp_file.name))
temp_file.write('contents')
name = temp_file.name
self.assertEquals(un_gzip_stream(gzip_file(name)).read(), 'contents')
def _stream_file(self, uuid, path, gzipped):
"""
Returns a file-like object reading the given file. This file is gzipped
if gzipped is True.
"""
if self._is_available_locally(uuid):
file_path = self._get_target_path(uuid, path)
if gzipped:
return file_util.gzip_file(file_path)
else:
return open(file_path)
else:
worker = self.get_bundle_worker(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, uuid, path,
read_args)
fileobj = self._get_read_response_stream(response_socket_id)
if not gzipped:
fileobj = file_util.un_gzip_stream(fileobj)
return Deallocating(fileobj, self._worker_model,
response_socket_id)
except:
self._worker_model.deallocate_socket(response_socket_id)
raise
def stream_file(self, uuid, path, gzipped):
"""
Returns a file-like object reading the given file. This file is gzipped
if gzipped is True.
"""
if self._is_available_locally(uuid):
file_path = self._get_target_path(uuid, path)
if gzipped:
return file_util.gzip_file(file_path)
else:
return open(file_path)
else:
worker = self._worker_model.get_bundle_worker(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, uuid, path, read_args)
fileobj = self._get_read_response_stream(response_socket_id)
if not gzipped:
fileobj = file_util.un_gzip_stream(fileobj)
return Deallocating(fileobj, self._worker_model, response_socket_id)
except Exception:
self._worker_model.deallocate_socket(response_socket_id)
raise
def unpack(ext, source, dest_path):
"""
Unpack the archive |source| to |dest_path|.
Note: |source| can be a file handle or a path.
|ext| contains the extension of the archive.
"""
if ext != '.zip':
close_source = False
try:
if isinstance(source, basestring):
source = open(source, 'rb')
close_source = True
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)
else:
raise UsageError('Not an archive.')
except (tarfile.TarError, IOError):
raise UsageError('Invalid archive upload.')
finally:
if close_source:
source.close()
else:
delete_source = False
try:
# unzip doesn't accept input from standard input, so we have to save
# to a temporary file.
if not isinstance(source, basestring):
temp_path = dest_path + '.zip'
with open(temp_path, 'wb') as f:
shutil.copyfileobj(source, f)
source = temp_path
delete_source = True
exitcode = subprocess.call(
['unzip', '-q', source, '-d', dest_path])
if exitcode != 0:
raise UsageError('Invalid archive upload.')
finally:
if delete_source:
path_util.remove(source)
def unpack(ext, source, dest_path):
"""
Unpack the archive |source| to |dest_path|.
Note: |source| can be a file handle or a path.
|ext| contains the extension of the archive.
"""
if ext != '.zip':
close_source = False
try:
if isinstance(source, basestring):
source = open(source, 'rb')
close_source = True
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)
else:
raise UsageError('Not an archive.')
except (tarfile.TarError, IOError):
raise UsageError('Invalid archive upload.')
finally:
if close_source:
source.close()
else:
delete_source = False
try:
# unzip doesn't accept input from standard input, so we have to save
# to a temporary file.
if not isinstance(source, basestring):
temp_path = dest_path + '.zip'
with open(temp_path, 'wb') as f:
shutil.copyfileobj(source, f)
source = temp_path
delete_source = True
exitcode = subprocess.call(['unzip', '-q', source, '-d', dest_path])
if exitcode != 0:
raise UsageError('Invalid archive upload.')
finally:
if delete_source:
path_util.remove(source)
