extract_archive can now extract zip format.
This commit is contained in:
parent
00acba9ee9
commit
f2c841df6d
50
utils.py
50
utils.py
|
@ -1,7 +1,8 @@
|
|||
import os.path
|
||||
import hashlib
|
||||
import tarfile
|
||||
import tarfile, zipfile
|
||||
import tempfile
|
||||
import os
|
||||
from collections import namedtuple, defaultdict
|
||||
|
||||
pj = os.path.join
|
||||
|
@ -63,29 +64,54 @@ class Context:
|
|||
|
||||
|
||||
def extract_archive(archive_path, dest_dir, topdir=None, name=None):
|
||||
with tarfile.open(archive_path) as archive:
|
||||
members = archive.getmembers()
|
||||
is_zip_archive = archive_path.endswith('.zip')
|
||||
try:
|
||||
if is_zip_archive:
|
||||
archive = zipfile.ZipFile(archive_path)
|
||||
members = archive.infolist()
|
||||
getname = lambda info : info.filename
|
||||
isdir = lambda info: info.filename.endswith('/')
|
||||
else:
|
||||
archive = tarfile.open(archive_path)
|
||||
members = archive.getmembers()
|
||||
getname = lambda info : getattr(info, 'name')
|
||||
isdir = lambda info: info.isdir()
|
||||
if not topdir:
|
||||
for d in (m for m in members if m.isdir()):
|
||||
if not os.path.dirname(d.name):
|
||||
for d in (m for m in members if isdir(m)):
|
||||
_name = getname(d)
|
||||
if _name.endswith('/'):
|
||||
_name = _name[:-1]
|
||||
if not os.path.dirname(_name):
|
||||
if topdir:
|
||||
# There is already a top dir.
|
||||
# Two topdirs in the same archive.
|
||||
# Extract all
|
||||
topdir = None
|
||||
break
|
||||
topdir = d
|
||||
else:
|
||||
topdir = archive.getmember(topdir)
|
||||
topdir = _name
|
||||
if topdir:
|
||||
members_to_extract = [m for m in members if m.name.startswith(topdir.name+'/')]
|
||||
members_to_extract = [m for m in members if getname(m).startswith(topdir+'/')]
|
||||
os.makedirs(dest_dir, exist_ok=True)
|
||||
with tempfile.TemporaryDirectory(prefix=os.path.basename(archive_path), dir=dest_dir) as tmpdir:
|
||||
archive.extractall(path=tmpdir, members=members_to_extract)
|
||||
name = name or topdir.name
|
||||
os.rename(pj(tmpdir, topdir.name), pj(dest_dir, name))
|
||||
if is_zip_archive:
|
||||
_members_to_extract = [getname(m) for m in members_to_extract]
|
||||
else:
|
||||
_members_to_extract = members_to_extract
|
||||
archive.extractall(path=tmpdir, members=_members_to_extract)
|
||||
if is_zip_archive:
|
||||
for member in members_to_extract:
|
||||
if isdir(member):
|
||||
continue
|
||||
perm = (member.external_attr >> 16) & 0x1FF
|
||||
os.chmod(pj(tmpdir, getname(member)), perm)
|
||||
name = name or topdir
|
||||
os.rename(pj(tmpdir, topdir), pj(dest_dir, name))
|
||||
else:
|
||||
if name:
|
||||
dest_dir = pj(dest_dir, name)
|
||||
os.makedirs(dest_dir)
|
||||
archive.extractall(path=dest_dir)
|
||||
finally:
|
||||
if archive is not None:
|
||||
archive.close()
|
||||
|
||||
|
|
Loading…
Reference in New Issue