extract_archive can now extract zip format.

This commit is contained in:
Matthieu Gautier 2017-02-22 14:19:50 +01:00
parent 00acba9ee9
commit f2c841df6d
1 changed files with 38 additions and 12 deletions

View File

@ -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()