def test_unpack_registery(self):
formats = get_unpack_formats()
def _boo(filename, extract_dir, extra):
self.assertEqual(extra, 1)
self.assertEqual(filename, 'stuff.boo')
self.assertEqual(extract_dir, 'xx')
register_unpack_format('Boo', ['.boo', '.b2'], _boo, [('extra', 1)])
unpack_archive('stuff.boo', 'xx')
# trying to register a .boo unpacker again
self.assertRaises(RegistryError, register_unpack_format, 'Boo2',
['.boo'], _boo)
# should work now
unregister_unpack_format('Boo')
register_unpack_format('Boo2', ['.boo'], _boo)
self.assertIn(('Boo2', ['.boo'], ''), get_unpack_formats())
self.assertNotIn(('Boo', ['.boo'], ''), get_unpack_formats())
# let's leave a clean state
unregister_unpack_format('Boo2')
self.assertEqual(get_unpack_formats(), formats)
def test_unpack_registery(self):
formats = get_unpack_formats()
def _boo(filename, extract_dir, extra):
self.assertEquals(extra, 1)
self.assertEquals(filename, 'stuff.boo')
self.assertEquals(extract_dir, 'xx')
register_unpack_format('Boo', ['.boo', '.b2'], _boo, [('extra', 1)])
unpack_archive('stuff.boo', 'xx')
# trying to register a .boo unpacker again
self.assertRaises(RegistryError, register_unpack_format, 'Boo2',
['.boo'], _boo)
# should work now
unregister_unpack_format('Boo')
register_unpack_format('Boo2', ['.boo'], _boo)
self.assertIn(('Boo2', ['.boo'], ''), get_unpack_formats())
self.assertNotIn(('Boo', ['.boo'], ''), get_unpack_formats())
# let's leave a clean state
unregister_unpack_format('Boo2')
self.assertEquals(get_unpack_formats(), formats)
def register():
"""
Configures shutil.unpack_archive to use our custom unzipper
"""
if sys.platform != "win32":
unregister_unpack_format('zip')
register_unpack_format('zip', ['.zip'],
_unpack_zipfile_with_permissions)
def prepare_shutil_state():
"""Reset any shutil modification in its current state"""
import shutil
registered_formats = [f[0] for f in shutil.get_unpack_formats()]
for format_id in tarball.ADDITIONAL_ARCHIVE_FORMATS:
name = format_id[0]
if name in registered_formats:
shutil.unregister_unpack_format(name)
return shutil
def unpack_archive(filename, extract_dir=None):
"""shutil.unpack_archive wrapper to unpack ['.dat.bz2'] archive.
:param filename: name of the archive.
:param extract_dir: name of the target directory, where the archive
is unpacked. If not provided, the current working directory is used.
"""
import shutil
# hardcoded for .dat.bz2
if filename.endswith('.dat.bz2'):
shutil.register_unpack_format(
'bzip2', ['dat.bz2'], _unpack_bz2, [], "bzip2'ed dat file")
shutil.unpack_archive(filename, extract_dir, 'bzip2')
shutil.unregister_unpack_format('bzip2')
else:
shutil.unpack_archive(filename, extract_dir)
def test_unpack_registery(self) -> None:
formats = get_unpack_formats()
def _boo(filename: str, extract_dir: str, extra: int) -> None:
self.assertEqual(extra, 1)
self.assertEqual(filename, "stuff.boo")
self.assertEqual(extract_dir, "xx")
register_unpack_format("Boo", [".boo", ".b2"], _boo, [("extra", 1)])
unpack_archive("stuff.boo", "xx")
# trying to register a .boo unpacker again
self.assertRaises(RegistryError, register_unpack_format, "Boo2", [".boo"], _boo)
# should work now
unregister_unpack_format("Boo")
register_unpack_format("Boo2", [".boo"], _boo)
self.assertIn(("Boo2", [".boo"], ""), get_unpack_formats())
self.assertNotIn(("Boo", [".boo"], ""), get_unpack_formats())
# let's leave a clean state
unregister_unpack_format("Boo2")
self.assertEqual(get_unpack_formats(), formats)
def unpack_gzip(archive_path: str, destination_dir: str):
destination_filepath = os.path.join(destination_dir,
Path(archive_path).stem)
with gzip.open(archive_path, 'rb') as f_in:
with open(destination_filepath, 'wb') as f_out:
copyfileobj(f_in, f_out)
def unpack_zip(archive_path: str, destination_dir: str):
zfile = zipfile.ZipFile(archive_path)
zfile.extractall(destination_dir)
register_unpack_format('gzip', ['.gz'], unpack_gzip)
unregister_unpack_format('zip')
register_unpack_format('zip', ['.zip'], unpack_zip)
@retry(delay=1, backoff=2, tries=4)
def download_and_extract(url: str,
destination_folder: str,
total_mb_size: Optional[float] = None) -> bool:
"""Download and extract from url, if file has already been downloaded return False else True."""
try:
# another process is currently (or was) working on the same url and destination folder
key = url + destination_folder
marker_filepath = os.path.join(
destination_folder,
sha256(key.encode()).hexdigest() + '.marker')
Path(destination_folder).mkdir(exist_ok=True)
def update_event(self, inp=-1):
self.set_output_val(0, shutil.unregister_unpack_format(self.input(0)))
