codekingpro/portable-devtools

1"""Utilities for extracting common archive formats"""2 3import contextlib4import os5import posixpath6import shutil7import tarfile8import zipfile9 10from ._path import ensure_directory11 12from distutils.errors import DistutilsError13 14__all__ = [15    "unpack_archive",16    "unpack_zipfile",17    "unpack_tarfile",18    "default_filter",19    "UnrecognizedFormat",20    "extraction_drivers",21    "unpack_directory",22]23 24 25class UnrecognizedFormat(DistutilsError):26    """Couldn't recognize the archive type"""27 28 29def default_filter(src, dst):30    """The default progress/filter callback; returns True for all files"""31    return dst32 33 34def unpack_archive(35    filename, extract_dir, progress_filter=default_filter, drivers=None36) -> None:37    """Unpack `filename` to `extract_dir`, or raise ``UnrecognizedFormat``38 39    `progress_filter` is a function taking two arguments: a source path40    internal to the archive ('/'-separated), and a filesystem path where it41    will be extracted.  The callback must return the desired extract path42    (which may be the same as the one passed in), or else ``None`` to skip43    that file or directory.  The callback can thus be used to report on the44    progress of the extraction, as well as to filter the items extracted or45    alter their extraction paths.46 47    `drivers`, if supplied, must be a non-empty sequence of functions with the48    same signature as this function (minus the `drivers` argument), that raise49    ``UnrecognizedFormat`` if they do not support extracting the designated50    archive type.  The `drivers` are tried in sequence until one is found that51    does not raise an error, or until all are exhausted (in which case52    ``UnrecognizedFormat`` is raised).  If you do not supply a sequence of53    drivers, the module's ``extraction_drivers`` constant will be used, which54    means that ``unpack_zipfile`` and ``unpack_tarfile`` will be tried, in that55    order.56    """57    for driver in drivers or extraction_drivers:58        try:59            driver(filename, extract_dir, progress_filter)60        except UnrecognizedFormat:61            continue62        else:63            return64    else:65        raise UnrecognizedFormat(f"Not a recognized archive type: {filename}")66 67 68def unpack_directory(filename, extract_dir, progress_filter=default_filter) -> None:69    """ "Unpack" a directory, using the same interface as for archives70 71    Raises ``UnrecognizedFormat`` if `filename` is not a directory72    """73    if not os.path.isdir(filename):74        raise UnrecognizedFormat(f"{filename} is not a directory")75 76    paths = {77        filename: ('', extract_dir),78    }79    for base, dirs, files in os.walk(filename):80        src, dst = paths[base]81        for d in dirs:82            paths[os.path.join(base, d)] = src + d + '/', os.path.join(dst, d)83        for f in files:84            target = os.path.join(dst, f)85            target = progress_filter(src + f, target)86            if not target:87                # skip non-files88                continue89            ensure_directory(target)90            f = os.path.join(base, f)91            shutil.copyfile(f, target)92            shutil.copystat(f, target)93 94 95def unpack_zipfile(filename, extract_dir, progress_filter=default_filter) -> None:96    """Unpack zip `filename` to `extract_dir`97 98    Raises ``UnrecognizedFormat`` if `filename` is not a zipfile (as determined99    by ``zipfile.is_zipfile()``).  See ``unpack_archive()`` for an explanation100    of the `progress_filter` argument.101    """102 103    if not zipfile.is_zipfile(filename):104        raise UnrecognizedFormat(f"{filename} is not a zip file")105 106    with zipfile.ZipFile(filename) as z:107        _unpack_zipfile_obj(z, extract_dir, progress_filter)108 109 110def _unpack_zipfile_obj(zipfile_obj, extract_dir, progress_filter=default_filter):111    """Internal/private API used by other parts of setuptools.112    Similar to ``unpack_zipfile``, but receives an already opened :obj:`zipfile.ZipFile`113    object instead of a filename.114    """115    for info in zipfile_obj.infolist():116        name = info.filename117 118        # don't extract absolute paths or ones with .. in them119        if name.startswith('/') or '..' in name.split('/'):120            continue121 122        target = os.path.join(extract_dir, *name.split('/'))123        target = progress_filter(name, target)124        if not target:125            continue126        if name.endswith('/'):127            # directory128            ensure_directory(target)129        else:130            # file131            ensure_directory(target)132            data = zipfile_obj.read(info.filename)133            with open(target, 'wb') as f:134                f.write(data)135        unix_attributes = info.external_attr >> 16136        if unix_attributes:137            os.chmod(target, unix_attributes)138 139 140def _resolve_tar_file_or_dir(tar_obj, tar_member_obj):141    """Resolve any links and extract link targets as normal files."""142    while tar_member_obj is not None and (143        tar_member_obj.islnk() or tar_member_obj.issym()144    ):145        linkpath = tar_member_obj.linkname146        if tar_member_obj.issym():147            base = posixpath.dirname(tar_member_obj.name)148            linkpath = posixpath.join(base, linkpath)149            linkpath = posixpath.normpath(linkpath)150        tar_member_obj = tar_obj._getmember(linkpath)151 152    is_file_or_dir = tar_member_obj is not None and (153        tar_member_obj.isfile() or tar_member_obj.isdir()154    )155    if is_file_or_dir:156        return tar_member_obj157 158    raise LookupError('Got unknown file type')159 160 161def _iter_open_tar(tar_obj, extract_dir, progress_filter):162    """Emit member-destination pairs from a tar archive."""163    # don't do any chowning!164    tar_obj.chown = lambda *args: None165 166    with contextlib.closing(tar_obj):167        for member in tar_obj:168            name = member.name169            # don't extract absolute paths or ones with .. in them170            if name.startswith('/') or '..' in name.split('/'):171                continue172 173            prelim_dst = os.path.join(extract_dir, *name.split('/'))174 175            try:176                member = _resolve_tar_file_or_dir(tar_obj, member)177            except LookupError:178                continue179 180            final_dst = progress_filter(name, prelim_dst)181            if not final_dst:182                continue183 184            if final_dst.endswith(os.sep):185                final_dst = final_dst[:-1]186 187            yield member, final_dst188 189 190def unpack_tarfile(filename, extract_dir, progress_filter=default_filter) -> bool:191    """Unpack tar/tar.gz/tar.bz2 `filename` to `extract_dir`192 193    Raises ``UnrecognizedFormat`` if `filename` is not a tarfile (as determined194    by ``tarfile.open()``).  See ``unpack_archive()`` for an explanation195    of the `progress_filter` argument.196    """197    try:198        tarobj = tarfile.open(filename)199    except tarfile.TarError as e:200        raise UnrecognizedFormat(201            f"{filename} is not a compressed or uncompressed tar file"202        ) from e203 204    for member, final_dst in _iter_open_tar(205        tarobj,206        extract_dir,207        progress_filter,208    ):209        try:210            # XXX Ugh211            tarobj._extract_member(member, final_dst)212        except tarfile.ExtractError:213            # chown/chmod/mkfifo/mknode/makedev failed214            pass215 216    return True217 218 219extraction_drivers = unpack_directory, unpack_zipfile, unpack_tarfile220