Aluode/PerceptionLabPortable

1"""distutils.archive_util2 3Utility functions for creating archive files (tarballs, zip files,4that sort of thing)."""5 6from __future__ import annotations7 8import os9from typing import Literal, overload10 11try:12    import zipfile13except ImportError:14    zipfile = None15 16 17from ._log import log18from .dir_util import mkpath19from .errors import DistutilsExecError20from .spawn import spawn21 22try:23    from pwd import getpwnam24except ImportError:25    getpwnam = None26 27try:28    from grp import getgrnam29except ImportError:30    getgrnam = None31 32 33def _get_gid(name):34    """Returns a gid, given a group name."""35    if getgrnam is None or name is None:36        return None37    try:38        result = getgrnam(name)39    except KeyError:40        result = None41    if result is not None:42        return result[2]43    return None44 45 46def _get_uid(name):47    """Returns an uid, given a user name."""48    if getpwnam is None or name is None:49        return None50    try:51        result = getpwnam(name)52    except KeyError:53        result = None54    if result is not None:55        return result[2]56    return None57 58 59def make_tarball(60    base_name: str,61    base_dir: str | os.PathLike[str],62    compress: Literal["gzip", "bzip2", "xz"] | None = "gzip",63    verbose: bool = False,64    dry_run: bool = False,65    owner: str | None = None,66    group: str | None = None,67) -> str:68    """Create a (possibly compressed) tar file from all the files under69    'base_dir'.70 71    'compress' must be "gzip" (the default), "bzip2", "xz", or None.72 73    'owner' and 'group' can be used to define an owner and a group for the74    archive that is being built. If not provided, the current owner and group75    will be used.76 77    The output tar file will be named 'base_dir' +  ".tar", possibly plus78    the appropriate compression extension (".gz", ".bz2", ".xz" or ".Z").79 80    Returns the output filename.81    """82    tar_compression = {83        'gzip': 'gz',84        'bzip2': 'bz2',85        'xz': 'xz',86        None: '',87    }88    compress_ext = {'gzip': '.gz', 'bzip2': '.bz2', 'xz': '.xz'}89 90    # flags for compression program, each element of list will be an argument91    if compress is not None and compress not in compress_ext.keys():92        raise ValueError(93            "bad value for 'compress': must be None, 'gzip', 'bzip2', 'xz'"94        )95 96    archive_name = base_name + '.tar'97    archive_name += compress_ext.get(compress, '')98 99    mkpath(os.path.dirname(archive_name), dry_run=dry_run)100 101    # creating the tarball102    import tarfile  # late import so Python build itself doesn't break103 104    log.info('Creating tar archive')105 106    uid = _get_uid(owner)107    gid = _get_gid(group)108 109    def _set_uid_gid(tarinfo):110        if gid is not None:111            tarinfo.gid = gid112            tarinfo.gname = group113        if uid is not None:114            tarinfo.uid = uid115            tarinfo.uname = owner116        return tarinfo117 118    if not dry_run:119        tar = tarfile.open(archive_name, f'w|{tar_compression[compress]}')120        try:121            tar.add(base_dir, filter=_set_uid_gid)122        finally:123            tar.close()124 125    return archive_name126 127 128def make_zipfile(  # noqa: C901129    base_name: str,130    base_dir: str | os.PathLike[str],131    verbose: bool = False,132    dry_run: bool = False,133) -> str:134    """Create a zip file from all the files under 'base_dir'.135 136    The output zip file will be named 'base_name' + ".zip".  Uses either the137    "zipfile" Python module (if available) or the InfoZIP "zip" utility138    (if installed and found on the default search path).  If neither tool is139    available, raises DistutilsExecError.  Returns the name of the output zip140    file.141    """142    zip_filename = base_name + ".zip"143    mkpath(os.path.dirname(zip_filename), dry_run=dry_run)144 145    # If zipfile module is not available, try spawning an external146    # 'zip' command.147    if zipfile is None:148        if verbose:149            zipoptions = "-r"150        else:151            zipoptions = "-rq"152 153        try:154            spawn(["zip", zipoptions, zip_filename, base_dir], dry_run=dry_run)155        except DistutilsExecError:156            # XXX really should distinguish between "couldn't find157            # external 'zip' command" and "zip failed".158            raise DistutilsExecError(159                f"unable to create zip file '{zip_filename}': "160                "could neither import the 'zipfile' module nor "161                "find a standalone zip utility"162            )163 164    else:165        log.info("creating '%s' and adding '%s' to it", zip_filename, base_dir)166 167        if not dry_run:168            try:169                zip = zipfile.ZipFile(170                    zip_filename, "w", compression=zipfile.ZIP_DEFLATED171                )172            except RuntimeError:173                zip = zipfile.ZipFile(zip_filename, "w", compression=zipfile.ZIP_STORED)174 175            with zip:176                if base_dir != os.curdir:177                    path = os.path.normpath(os.path.join(base_dir, ''))178                    zip.write(path, path)179                    log.info("adding '%s'", path)180                for dirpath, dirnames, filenames in os.walk(base_dir):181                    for name in dirnames:182                        path = os.path.normpath(os.path.join(dirpath, name, ''))183                        zip.write(path, path)184                        log.info("adding '%s'", path)185                    for name in filenames:186                        path = os.path.normpath(os.path.join(dirpath, name))187                        if os.path.isfile(path):188                            zip.write(path, path)189                            log.info("adding '%s'", path)190 191    return zip_filename192 193 194ARCHIVE_FORMATS = {195    'gztar': (make_tarball, [('compress', 'gzip')], "gzip'ed tar-file"),196    'bztar': (make_tarball, [('compress', 'bzip2')], "bzip2'ed tar-file"),197    'xztar': (make_tarball, [('compress', 'xz')], "xz'ed tar-file"),198    'ztar': (make_tarball, [('compress', 'compress')], "compressed tar file"),199    'tar': (make_tarball, [('compress', None)], "uncompressed tar file"),200    'zip': (make_zipfile, [], "ZIP file"),201}202 203 204def check_archive_formats(formats):205    """Returns the first format from the 'format' list that is unknown.206 207    If all formats are known, returns None208    """209    for format in formats:210        if format not in ARCHIVE_FORMATS:211            return format212    return None213 214 215@overload216def make_archive(217    base_name: str,218    format: str,219    root_dir: str | os.PathLike[str] | bytes | os.PathLike[bytes] | None = None,220    base_dir: str | None = None,221    verbose: bool = False,222    dry_run: bool = False,223    owner: str | None = None,224    group: str | None = None,225) -> str: ...226@overload227def make_archive(228    base_name: str | os.PathLike[str],229    format: str,230    root_dir: str | os.PathLike[str] | bytes | os.PathLike[bytes],231    base_dir: str | None = None,232    verbose: bool = False,233    dry_run: bool = False,234    owner: str | None = None,235    group: str | None = None,236) -> str: ...237def make_archive(238    base_name: str | os.PathLike[str],239    format: str,240    root_dir: str | os.PathLike[str] | bytes | os.PathLike[bytes] | None = None,241    base_dir: str | None = None,242    verbose: bool = False,243    dry_run: bool = False,244    owner: str | None = None,245    group: str | None = None,246) -> str:247    """Create an archive file (eg. zip or tar).248 249    'base_name' is the name of the file to create, minus any format-specific250    extension; 'format' is the archive format: one of "zip", "tar", "gztar",251    "bztar", "xztar", or "ztar".252 253    'root_dir' is a directory that will be the root directory of the254    archive; ie. we typically chdir into 'root_dir' before creating the255    archive.  'base_dir' is the directory where we start archiving from;256    ie. 'base_dir' will be the common prefix of all files and257    directories in the archive.  'root_dir' and 'base_dir' both default258    to the current directory.  Returns the name of the archive file.259 260    'owner' and 'group' are used when creating a tar archive. By default,261    uses the current owner and group.262    """263    save_cwd = os.getcwd()264    if root_dir is not None:265        log.debug("changing into '%s'", root_dir)266        base_name = os.path.abspath(base_name)267        if not dry_run:268            os.chdir(root_dir)269 270    if base_dir is None:271        base_dir = os.curdir272 273    kwargs = {'dry_run': dry_run}274 275    try:276        format_info = ARCHIVE_FORMATS[format]277    except KeyError:278        raise ValueError(f"unknown archive format '{format}'")279 280    func = format_info[0]281    kwargs.update(format_info[1])282 283    if format != 'zip':284        kwargs['owner'] = owner285        kwargs['group'] = group286 287    try:288        filename = func(base_name, base_dir, **kwargs)289    finally:290        if root_dir is not None:291            log.debug("changing back to '%s'", save_cwd)292            os.chdir(save_cwd)293 294    return filename295