python3-pyutilib 5.3.5-1 / usr / lib / python3 / dist-packages / pyutilib / misc / archivereader.py

#  _________________________________________________________________________
#
#  PyUtilib: A Python utility library.
#  Copyright (c) 2008 Sandia Corporation.
#  This software is distributed under the BSD License.
#  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
#  the U.S. Government retains certain rights in this software.
#  _________________________________________________________________________

__all__ = ['ArchiveReaderFactory', 'ArchiveReader',
           'ZipArchiveReader', 'TarArchiveReader',
           'DirArchiveReader', 'FileArchiveReader',
           'GzipFileArchiveReader', 'BZ2FileArchiveReader']

import os
import sys
import tempfile
import shutil
import posixpath

_sep = '/'

zipfile_available = False
tarfile_available = False
gzip_available = False
bz2_available = False
if sys.version_info[:2] >= (2,6):
    try:
        import zipfile
        zipfile_available = True
    except:
        pass
try:
    import tarfile
    tarfile_available = True
except:
    pass
try:
    import gzip
    gzip_available = True
except:
    pass
try:
    import bz2
    bz2_available = True
except:
    pass

_WindowsError = None
try:
    _WindowsError = WindowsError
except NameError:
    # not windows
    pass

#
# This class presents a simple interface for unpacking
# archives such a .zip and .tar (bz2 or gzip) as well as
# normal directories. Keywords make convenient it to partially unpack
# archives based on a subdirectory name or maximum recursion depth,
#
#
#
# Bugs:
#  * Python 2.5 and 2.6: The tarfile module incorrectly recognizes plain
#    text files as tar archives. This prevents an exception from being thrown
#    in the ArchiveReaderFactory when a non-archive/directory element is provided.
#

def ArchiveReaderFactory(dirname, **kwds):
    if not os.path.exists(ArchiveReader.normalize_name(dirname)):
        raise IOError("Cannot find file or directory `"
                      +dirname+"'\nPath expanded to: '"+
                      ArchiveReader.normalize_name(dirname)+"'")
    if ArchiveReader.isDir(dirname):
        return DirArchiveReader(dirname, **kwds)
    elif zipfile_available and ArchiveReader.isZip(dirname):
        return ZipArchiveReader(dirname, **kwds)
    elif tarfile_available and ArchiveReader.isTar(dirname):
        return TarArchiveReader(dirname, **kwds)
    elif gzip_available and ArchiveReader.isGzipFile(dirname):
        return GzipFileArchiveReader(dirname, **kwds)
    elif bz2_available and ArchiveReader.isBZ2File(dirname):
        return BZ2FileArchiveReader(dirname, **kwds)
    elif ArchiveReader.isFile(dirname):
        return FileArchiveReader(dirname, **kwds)
    else:
        raise ValueError("ArchiveReaderFactory was given an "
                         "unrecognized archive type with "
                         "name '%s'" % dirname)

class ArchiveReader(object):

    @staticmethod
    def isDir(name):
        return os.path.isdir(ArchiveReader.normalize_name(name))

    @staticmethod
    def isZip(name):
        if not zipfile_available:
            raise ImportError("zipfile support is disabled")
        return zipfile.is_zipfile(ArchiveReader.normalize_name(name))

    @staticmethod
    def isTar(name):
        if not tarfile_available:
            raise ImportError("tarfile support is disabled")
        return tarfile.is_tarfile(ArchiveReader.normalize_name(name))

    @staticmethod
    def isArchivedFile(name):
        return (not (ArchiveReader.isDir(name) or ArchiveReader.isFile(name))) and \
            (tarfile_available and ArchiveReader.isTar(name)) or \
            (zipfile_available and ArchiveReader.isZip(name)) or \
            (gzip_available and ArchiveReader.isGzipFile(name)) or \
            (bz2_available and ArchiveReader.isBZ2File(name))

    @staticmethod
    def isGzipFile(name):
        if not gzip_available:
            raise ImportError("gzip support is disabled")
        try:
            f = gzip.GzipFile(ArchiveReader.normalize_name(name))
            f.read(1)
            f.close()
        except IOError:
            return False
        return True

    @staticmethod
    def isBZ2File(name):
        if not bz2_available:
            raise ImportError("bz2 support is disabled")
        try:
            f = bz2.BZ2File(ArchiveReader.normalize_name(name))
            f.read(1)
            f.close()
        except IOError:
            return False
        return True

    @staticmethod
    def isFile(name):
        return os.path.isfile(ArchiveReader.normalize_name(name))

    @staticmethod
    def normalize_name(filename):
        """Turns the given file name into a normalized absolute path"""
        if filename is not None:
            filename = os.path.expanduser(filename)
            if not os.path.isabs(filename):
                filename = os.path.abspath(filename)
            filename = ArchiveReader._posix_name(filename)
            return posixpath.normpath(filename)

    @staticmethod
    def _posix_name(filename):
        if filename is not None:
            return filename.replace('\\','/')

    def __init__(self, name, *args, **kwds):
        posixabsname = self.normalize_name(name)
        if not os.path.exists(posixabsname):
            raise IOError("cannot find file or directory `"
                          +posixabsname+"'")

        self._abspath = os.path.dirname(posixabsname)
        self._basename = os.path.basename(posixabsname)
        self._archive_name = posixabsname

        subdir = kwds.pop('subdir', None)
        if (subdir is not None) and (subdir.strip() == ''):
            subdir = None
        maxdepth = kwds.pop('maxdepth', None)
        self._filter = kwds.pop('filter',None)

        self._subdir = posixpath.normpath(ArchiveReader._posix_name(subdir))+_sep \
                       if (subdir is not None) else None

        self._maxdepth = maxdepth
        if (self._maxdepth is not None) and (self._maxdepth < 0):
            raise ValueError("maxdepth must be >= 0")

        self._names_list = []
        self._artificial_dirs = set()
        self._extractions = set()
        # the python zipfile or tarfile object or None for (dir)
        self._handler = None
        self._workdir = tempfile.mkdtemp()

        self._init(*args, **kwds)

        if self._filter is not None:
            self._names_list = [_f for _f in self._names_list \
                                if not self._filter(_f)]

    def name(self):
        return self._archive_name

    def getExtractionDir(self):
        return self._workdir

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.close()

    def close(self):
        if self._handler is not None:
            self._handler.close()
        if (self._workdir is not None) and \
           (os.path.exists(self._workdir)):
            shutil.rmtree(self._workdir, True)
        self._workdir = None

    def clear_extractions(self):
        # don't do anything to directories since they may have existed
        # prior to extracting them and we don't know what they
        # contained
        for name in self._extractions:
            if os.path.exists(name) and \
               (not os.path.isdir(name)):
                os.remove(name)
        self._extractions.clear()
        # If the directories existed in the workdir, then they will be
        # taken care of here
        if (self._workdir is not None) and \
           (os.path.exists(self._workdir)):
            for root, dirs, files in os.walk(self._workdir):
                while len(dirs):
                    shutil.rmtree(posixpath.join(root,dirs.pop()), True)
                while len(files):
                    os.remove(posixpath.join(root,files.pop()))

    def getnames(self):
        return self._names_list

    def contains(self, name):
        return name in self._names_list

    def _validate_name(self, name):
        name = self._posix_name(name)
        if (name is None) and (len(self._names_list) == 0):
            raise KeyError("The archive is empty!")
        elif (name is None) and (len(self._names_list) > 1):
            raise KeyError("A name argument is required when "
                           "the archive has more than one member")
        elif name is None:
            name = self._names_list[0]

        if name in self._artificial_dirs:
            return None, name

        if not self.contains(name):
            #
            # There are cases (e.g., Windows zipped folders) where
            # directory names don't appear in the list of objects
            # available to in zipfile. We will try to retroactively
            # handle this case by check if any listed archive members
            # begin with name+_sep, and if so, adding the directory
            # name to the _names_list
            #
            checkname = name+_sep
            if (self._maxdepth is None) or (checkname.count(_sep) <= self._maxdepth+1):
                for othername in self._fulldepth_names_list:
                    if othername.startswith(checkname):
                        self._artificial_dirs.add(name)
                        self._names_list.append(name)
                        return None, name
            msg = ("There is no item named '%s' in "
                   "the archive %s" % (name, self._basename))
            if self._subdir is not None:
                msg += ", subdirectory: "+self._subdir
            raise KeyError(msg)
        absname = name if (self._subdir is None) else self._subdir+name
        return absname, name

    def open(self, name=None, *args, **kwds):
        absname, relname = self._validate_name(name)
        return self._openImp(absname, relname, *args, **kwds)

    def _openImp(self, name, *args, **kwds):
        raise NotImplementedError("This method has not been "
                                  "implemented")

    def extract(self, member=None, path=None, recursive=False, *args, **kwds):
        absolute_name, relative_name = self._validate_name(member)
        dst, children = self._extractImp(absolute_name,
                                         relative_name,
                                         path,
                                         recursive,
                                         *args,
                                         **kwds)
        self._extractions.add(dst)
        if not recursive:
            assert len(children) == 0
        else:
            self._extractions.update(children)
        return dst

    def _extractImp(self, absolute_name, relative_name, path, recursive, *args, **kwds):
        raise NotImplementedError("This method has not been "
                                  "implemented")

    def extractall(self, path=None,  members=None, recursive=False, *args, **kwds):
        names = None
        if members is not None:
            names = set(members)
        else:
            names = set(self._names_list)
        # Save the expense checking recursions if there is no point
        if len(names) == len(self._names_list):
            recursive = False
        dsts = []
        while names:
            absolute_name, relative_name = self._validate_name(names.pop())
            dst, children = self._extractImp(absolute_name,
                                             relative_name,
                                             path,
                                             recursive,
                                             *args,
                                             **kwds)
            if not recursive:
                assert len(children) == 0
            dsts.append(dst)
            self._extractions.add(dst)
            if len(children):
                self._extractions.update(children)
                names = names-set(children)

        return dsts

    # like shutil.copytree, but will handle existing
    # directories the same way tarfile extraction
    # occurs by adding new content rather than
    # raising an exception
    @staticmethod
    def _copytree(src, dst, ignores=None, maxdepth=None):

        assert os.path.exists(src) and os.path.isdir(src)
        if not os.path.exists(dst):
            os.makedirs(dst)

        if (maxdepth is None) or (maxdepth > 0):

            if maxdepth is not None:
                maxdepth -= 1

            names = os.listdir(src)
            if ignores is not None:
                ignored_names = ignores
            else:
                ignored_names = set()

            for name in names:
                srcname = posixpath.join(src, name)
                if srcname in ignored_names:
                    continue
                dstname = posixpath.join(dst, name)
                if os.path.isdir(srcname):
                    ArchiveReader._copytree(srcname, dstname, ignores=ignores, maxdepth=maxdepth)
                else:
                    shutil.copy2(srcname, dstname)
        try:
            shutil.copystat(src, dst)
        except _WindowsError:
            # can't copy file access times on Windows
            pass

class _ziptar_base(ArchiveReader):

    @staticmethod
    def _fixnames(names, subdir, maxdepth):
        names_list = [posixpath.normpath(name) for name in names]
        subdir_depth = 0
        if subdir is not None:
            names_list = [name for name in names_list \
                          if name.startswith(subdir)]
            subdir_depth = subdir.count(_sep)
        fulldepth_names_list = list(names_list)
        if maxdepth is not None:
            names_list = [name for name in fulldepth_names_list \
                          if (name.count(_sep) - subdir_depth) <= maxdepth]
        if subdir is not None:
            names_list = [name.replace(subdir,'') \
                          for name in names_list]
            fulldepth_names_list = [name.replace(subdir,'') \
                               for name in fulldepth_names_list]
        return names_list, fulldepth_names_list, subdir_depth

    def _extractImp(self, absolute_name, relative_name, path, recursive):

        use_handler = True
        if absolute_name is None:
            # This case implies that this was an artificially
            # added directory that was not appearing in the
            # archive even though it technically exists
            # (a rare but possible edge case)
            use_handler = False
            absolute_name = relative_name if (self._subdir is None) else self._subdir+relative_name

        tmp_dst = posixpath.join(self._workdir, absolute_name)

        if use_handler:
            try:
                self._handler.extract(absolute_name, self._workdir)
            except KeyError: # sometimes directories need an _sep ending
                self._handler.extract(absolute_name+_sep, self._workdir)
        else:
            if not os.path.exists(tmp_dst):
                os.makedirs(tmp_dst)

        dst = tmp_dst
        if path is not None:

            dst = posixpath.join(path, relative_name)
            if os.path.isdir(tmp_dst):
                # updated the timestamp if it exists, otherwise
                # just create the directory
                self._copytree(tmp_dst, dst, maxdepth=0)
            else:
                if not os.path.exists(os.path.dirname(dst)):
                    os.makedirs(os.path.dirname(dst))
                try:
                    os.rename(tmp_dst, dst)
                except OSError:
                    shutil.copy2(tmp_dst, dst)

        children = []
        if os.path.isdir(dst) and recursive:
            new_names = [relname for relname in self._names_list \
                         if relname.startswith(relative_name)]
            if relative_name in new_names:
                new_names.remove(relative_name)
            for childname in new_names:
                absolute_childname, relative_childname = self._validate_name(childname)
                childdst, recursives = self._extractImp(absolute_childname, relative_childname, path, False)
                assert len(recursives) == 0
                children.append(childdst)

        return dst, children

class ZipArchiveReader(_ziptar_base):

    def _init(self, *args, **kwds):
        if zipfile is None:
            raise ImportError("zipfile support is disabled")
        assert(self._abspath is not None)
        assert(self._basename is not None)
        assert(self._archive_name is not None)

        if not self.isZip(self._archive_name):
            raise TypeError("Unrecognized zipfile format for file: %s"
                            % (self._archive_name))

        self._handler = zipfile.ZipFile(self._archive_name, *args, **kwds)
        self._names_list, self._fulldepth_names_list, self._subdir_depth = \
            self._fixnames(self._handler.namelist(), self._subdir, self._maxdepth)

    def _openImp(self, absolute_name, relative_name, *args, **kwds):
        f = None
        try:
            if absolute_name is None:
                raise KeyError
            f = self._handler.open(absolute_name, *args, **kwds)
        except KeyError:
            # when this method is called we have already verified the name
            # existed in the list, so this must be a directory
            raise IOError("Failed to open with zipfile, this must be a "
                          "directory: %s" % (absolute_name))
        return f

class TarArchiveReader(_ziptar_base):

    def _init(self, *args, **kwds):
        if not tarfile_available:
            raise ImportError("tarfile support is disabled")
        assert(self._abspath is not None)
        assert(self._basename is not None)
        assert(self._archive_name is not None)

        if not self.isTar(self._archive_name):
            raise TypeError("Unrecognized tarfile format for file: %s"
                            % (self._archive_name))

        self._handler = tarfile.open(self._archive_name, *args, **kwds)
        self._names_list, self._fulldepth_names_list, self._subdir_depth = \
            self._fixnames(self._handler.getnames(), self._subdir, self._maxdepth)

    def _openImp(self, absolute_name, relative_name, *args, **kwds):
        f = None
        if absolute_name is not None:
            f = self._handler.extractfile(absolute_name, *args, **kwds)
        if f is None:
            # when this method is called we have already verified the name
            # existed in the list, so this must be a directory
            raise IOError("Failed to open with tarfile, this must be a "
                          "directory: %s" % (absolute_name))
        return f

class DirArchiveReader(ArchiveReader):

    def _init(self, *args, **kwds):
        assert(self._abspath is not None)
        assert(self._basename is not None)
        assert(self._archive_name is not None)

        if kwds:
            raise ValueError("Unexpected keyword options found "
                             "while initializing '%s':\n\t%s"
                             % ( type(self).__name__,
                                 ','.join(sorted(kwds.keys()))))
        if args:
            raise ValueError("Unexpected arguments found "
                             "while initializing '%s':\n\t%s"
                             % ( type(self).__name__,
                                 ','.join(args)))

        if not self.isDir(self._archive_name):
            raise TypeError("Path not recognized as a directory: %s"
                            % (self._archive_name))

        rootdir = self._archive_name
        if self._subdir is not None:
            rootdir = posixpath.join(rootdir,self._subdir)
            if not os.path.exists(rootdir):
                raise IOError("Subdirectory '%s' does not exists in root directory: %s"
                              % (self._subdir, self._archive_name))
            self._names_list = self._walk(rootdir, maxdepth=self._maxdepth+1)
        else:
            self._names_list = self._walk(rootdir, maxdepth=self._maxdepth)
        self._fulldepth_names_list = self._walk(rootdir)

    @staticmethod
    def _walk(rootdir, maxdepth=None):
        names_list = []
        for root, dirs, files in os.walk(rootdir, topdown=True):
            prefix = posixpath.relpath(ArchiveReader._posix_name(root), rootdir)
            if prefix.endswith(_sep):
                prefix = prefix[:-1]
            if prefix == '.':
                prefix = ''
            for dname in dirs:
                names_list.append(posixpath.join(prefix,dname))
            if maxdepth is not None and prefix.count(_sep) >= maxdepth:
                continue
            for fname in files:
                names_list.append(posixpath.join(prefix,fname))
        return names_list

    def _extractImp(self, absolute_name, relative_name, path, recursive):

        assert absolute_name is not None

        if path is not None:
            dst = posixpath.join(path, relative_name)
        else:
            dst = posixpath.join(self._workdir, absolute_name)
        src = posixpath.join(self._archive_name, absolute_name)
        children = []
        if os.path.isdir(src):
            if recursive:
                ignores = []
                for child in self._walk(src):
                    rname = posixpath.join(relative_name, child)
                    if rname in self._names_list:
                        children.append(rname)
                    else:
                        ignores.append(posixpath.join(src, child))
                self._copytree(src, dst, ignores=ignores)

            elif not os.path.exists(dst):
                os.makedirs(dst)
            # if not recursive and the destination is an existing
            # directory, then there is nothing to do
        else:
            if not os.path.exists(posixpath.dirname(dst)):
                os.makedirs(os.path.dirname(dst))
            shutil.copy2(src, dst)

        return dst, children

    def _openImp(self, absolute_name, relative_name, *args, **kwds):
        assert absolute_name is not None
        return open(posixpath.join(self._archive_name, absolute_name), 'rb', *args, **kwds)

class FileArchiveReader(ArchiveReader):

    _handler_class = open

    def _extract_name(self,name):
        return name

    def _init(self):
        assert(self._abspath is not None)
        assert(self._basename is not None)
        assert(self._archive_name is not None)

        if self._subdir is not None:
            raise ValueError("'subdir' keyword option is not handled by "
                             "'%s'" % (type(self).__name__))
        if self._maxdepth is not None:
            raise ValueError("'maxdepth' keyword option is not handled by "
                             "'%s'" % (type(self).__name__))

        if not self.isFile(self._archive_name):
            raise TypeError("Path does not point to a file: %s"
                            % (self._archive_name))
        extract_name = self._extract_name(self._basename)
        if extract_name is not None:
            self._names_list = [extract_name]
            self._fulldepth_names_list = [extract_name]
            self._extract_name = extract_name
        else:
            self._names_list = [self._basename]
            self._fulldepth_names_list = [self._basename]
            self._extract_name = None

    def _extractImp(self, absolute_name, relative_name, path, recursive):
        assert absolute_name == self._extract_name
        assert relative_name == self._extract_name
        if recursive:
            raise ValueError("Recursive extraction does not make "
                             "sense for compressed file archive types")
        if path is not None:
            dst = posixpath.join(path, relative_name)
        else:
            dst = posixpath.join(self._workdir, absolute_name)
        with open(dst,'wb') as dstf:
            handler = self._handler_class(self._archive_name,'rb')
            shutil.copyfileobj(handler, dstf)
            handler.close()

        return dst, []

    def _openImp(self, absolute_name, relative_name, *args, **kwds):
        assert absolute_name == self._extract_name
        return self._handler_class(self._archive_name, 'rb', *args, **kwds)

class GzipFileArchiveReader(FileArchiveReader):

    _handler_class = gzip.GzipFile if gzip_available else None

    def _extract_name(self, name):
        # see the man page for gzip
        basename, ext = os.path.splitext(name)
        if ext == self._suffix:
            return basename
        else:
            return None

    def __init__(self, *args, **kwds):
        if not gzip_available:
            raise ImportError("gzip support is disabled")
        self._suffix = kwds.pop('suffix','.gz')
        super(GzipFileArchiveReader,self).__init__(*args, **kwds)
        if not self.isGzipFile(self._archive_name):
            raise TypeError("Unrecognized gzip format for file: %s"
                            % (self._archive_name))

    def _extractImp(self, absolute_name, relative_name, path, recursive):
        if self._extract_name is None:
            raise TypeError("Extraction disabled. File suffix %s does not "
                            "match expected suffix for compression type %s. "
                            "The default suffix can be changed by passing "
                            "'suffix=.<ext>' into the ArchiveReader constructor."
                            % (os.path.splitext(self._basename)[1],
                               self._suffix))
        return FileArchiveReader._extractImp(self,
                                             absolute_name,
                                             relative_name,
                                             path,
                                             recursive)

class BZ2FileArchiveReader(FileArchiveReader):

    _handler_class = bz2.BZ2File if bz2_available else None

    def _extract_name(self, name):
        # see the man page for bzip2
        basename, ext = os.path.splitext(name)
        if ext in ('.bz2','.bz'):
            return basename
        elif ext in ('.tbz2','.tbz'):
            return basename+'.tar'
        else:
            return name+'.out'

    def __init__(self, *args, **kwds):
        if not bz2_available:
            raise ImportError("bz2 support is disabled")
        super(BZ2FileArchiveReader,self).__init__(*args, **kwds)
        if not self.isBZ2File(self._archive_name):
            raise TypeError("Unrecognized bzip2 format for file: %s"
                            % (self._archive_name))