archive_util.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. """Utilities for extracting common archive formats"""
  2. import zipfile
  3. import tarfile
  4. import os
  5. import shutil
  6. import posixpath
  7. import contextlib
  8. from distutils.errors import DistutilsError
  9. from pkg_resources import ensure_directory
  10. __all__ = [
  11. "unpack_archive", "unpack_zipfile", "unpack_tarfile", "default_filter",
  12. "UnrecognizedFormat", "extraction_drivers", "unpack_directory",
  13. ]
  14. class UnrecognizedFormat(DistutilsError):
  15. """Couldn't recognize the archive type"""
  16. def default_filter(src, dst):
  17. """The default progress/filter callback; returns True for all files"""
  18. return dst
  19. def unpack_archive(
  20. filename, extract_dir, progress_filter=default_filter,
  21. drivers=None):
  22. """Unpack `filename` to `extract_dir`, or raise ``UnrecognizedFormat``
  23. `progress_filter` is a function taking two arguments: a source path
  24. internal to the archive ('/'-separated), and a filesystem path where it
  25. will be extracted. The callback must return the desired extract path
  26. (which may be the same as the one passed in), or else ``None`` to skip
  27. that file or directory. The callback can thus be used to report on the
  28. progress of the extraction, as well as to filter the items extracted or
  29. alter their extraction paths.
  30. `drivers`, if supplied, must be a non-empty sequence of functions with the
  31. same signature as this function (minus the `drivers` argument), that raise
  32. ``UnrecognizedFormat`` if they do not support extracting the designated
  33. archive type. The `drivers` are tried in sequence until one is found that
  34. does not raise an error, or until all are exhausted (in which case
  35. ``UnrecognizedFormat`` is raised). If you do not supply a sequence of
  36. drivers, the module's ``extraction_drivers`` constant will be used, which
  37. means that ``unpack_zipfile`` and ``unpack_tarfile`` will be tried, in that
  38. order.
  39. """
  40. for driver in drivers or extraction_drivers:
  41. try:
  42. driver(filename, extract_dir, progress_filter)
  43. except UnrecognizedFormat:
  44. continue
  45. else:
  46. return
  47. else:
  48. raise UnrecognizedFormat(
  49. "Not a recognized archive type: %s" % filename
  50. )
  51. def unpack_directory(filename, extract_dir, progress_filter=default_filter):
  52. """"Unpack" a directory, using the same interface as for archives
  53. Raises ``UnrecognizedFormat`` if `filename` is not a directory
  54. """
  55. if not os.path.isdir(filename):
  56. raise UnrecognizedFormat("%s is not a directory" % filename)
  57. paths = {
  58. filename: ('', extract_dir),
  59. }
  60. for base, dirs, files in os.walk(filename):
  61. src, dst = paths[base]
  62. for d in dirs:
  63. paths[os.path.join(base, d)] = src + d + '/', os.path.join(dst, d)
  64. for f in files:
  65. target = os.path.join(dst, f)
  66. target = progress_filter(src + f, target)
  67. if not target:
  68. # skip non-files
  69. continue
  70. ensure_directory(target)
  71. f = os.path.join(base, f)
  72. shutil.copyfile(f, target)
  73. shutil.copystat(f, target)
  74. def unpack_zipfile(filename, extract_dir, progress_filter=default_filter):
  75. """Unpack zip `filename` to `extract_dir`
  76. Raises ``UnrecognizedFormat`` if `filename` is not a zipfile (as determined
  77. by ``zipfile.is_zipfile()``). See ``unpack_archive()`` for an explanation
  78. of the `progress_filter` argument.
  79. """
  80. if not zipfile.is_zipfile(filename):
  81. raise UnrecognizedFormat("%s is not a zip file" % (filename,))
  82. with zipfile.ZipFile(filename) as z:
  83. for info in z.infolist():
  84. name = info.filename
  85. # don't extract absolute paths or ones with .. in them
  86. if name.startswith('/') or '..' in name.split('/'):
  87. continue
  88. target = os.path.join(extract_dir, *name.split('/'))
  89. target = progress_filter(name, target)
  90. if not target:
  91. continue
  92. if name.endswith('/'):
  93. # directory
  94. ensure_directory(target)
  95. else:
  96. # file
  97. ensure_directory(target)
  98. data = z.read(info.filename)
  99. with open(target, 'wb') as f:
  100. f.write(data)
  101. unix_attributes = info.external_attr >> 16
  102. if unix_attributes:
  103. os.chmod(target, unix_attributes)
  104. def unpack_tarfile(filename, extract_dir, progress_filter=default_filter):
  105. """Unpack tar/tar.gz/tar.bz2 `filename` to `extract_dir`
  106. Raises ``UnrecognizedFormat`` if `filename` is not a tarfile (as determined
  107. by ``tarfile.open()``). See ``unpack_archive()`` for an explanation
  108. of the `progress_filter` argument.
  109. """
  110. try:
  111. tarobj = tarfile.open(filename)
  112. except tarfile.TarError:
  113. raise UnrecognizedFormat(
  114. "%s is not a compressed or uncompressed tar file" % (filename,)
  115. )
  116. with contextlib.closing(tarobj):
  117. # don't do any chowning!
  118. tarobj.chown = lambda *args: None
  119. for member in tarobj:
  120. name = member.name
  121. # don't extract absolute paths or ones with .. in them
  122. if not name.startswith('/') and '..' not in name.split('/'):
  123. prelim_dst = os.path.join(extract_dir, *name.split('/'))
  124. # resolve any links and to extract the link targets as normal
  125. # files
  126. while member is not None and (
  127. member.islnk() or member.issym()):
  128. linkpath = member.linkname
  129. if member.issym():
  130. base = posixpath.dirname(member.name)
  131. linkpath = posixpath.join(base, linkpath)
  132. linkpath = posixpath.normpath(linkpath)
  133. member = tarobj._getmember(linkpath)
  134. if member is not None and (member.isfile() or member.isdir()):
  135. final_dst = progress_filter(name, prelim_dst)
  136. if final_dst:
  137. if final_dst.endswith(os.sep):
  138. final_dst = final_dst[:-1]
  139. try:
  140. # XXX Ugh
  141. tarobj._extract_member(member, final_dst)
  142. except tarfile.ExtractError:
  143. # chown/chmod/mkfifo/mknode/makedev failed
  144. pass
  145. return True
  146. extraction_drivers = unpack_directory, unpack_zipfile, unpack_tarfile