diff options
Diffstat (limited to 'testing/mozharness/mozfile/mozfile.py')
| -rw-r--r-- | testing/mozharness/mozfile/mozfile.py | 372 |
1 files changed, 372 insertions, 0 deletions
diff --git a/testing/mozharness/mozfile/mozfile.py b/testing/mozharness/mozfile/mozfile.py new file mode 100644 index 000000000..ac0edcab4 --- /dev/null +++ b/testing/mozharness/mozfile/mozfile.py @@ -0,0 +1,372 @@ +# -*- coding: utf-8 -*- + +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +from contextlib import contextmanager +import os +import shutil +import stat +import tarfile +import tempfile +import urlparse +import urllib2 +import zipfile +import time + +__all__ = ['extract_tarball', + 'extract_zip', + 'extract', + 'is_url', + 'load', + 'remove', + 'rmtree', + 'tree', + 'NamedTemporaryFile', + 'TemporaryDirectory'] + +try: + WindowsError +except NameError: + WindowsError = None # so we can unconditionally catch it later... + + +### utilities for extracting archives + +def extract_tarball(src, dest): + """extract a .tar file""" + + bundle = tarfile.open(src) + namelist = bundle.getnames() + + for name in namelist: + bundle.extract(name, path=dest) + bundle.close() + return namelist + + +def extract_zip(src, dest): + """extract a zip file""" + + if isinstance(src, zipfile.ZipFile): + bundle = src + else: + try: + bundle = zipfile.ZipFile(src) + except Exception, e: + print "src: %s" % src + raise + + namelist = bundle.namelist() + + for name in namelist: + filename = os.path.realpath(os.path.join(dest, name)) + if name.endswith('/'): + if not os.path.isdir(filename): + os.makedirs(filename) + else: + path = os.path.dirname(filename) + if not os.path.isdir(path): + os.makedirs(path) + _dest = open(filename, 'wb') + _dest.write(bundle.read(name)) + _dest.close() + mode = bundle.getinfo(name).external_attr >> 16 & 0x1FF + os.chmod(filename, mode) + bundle.close() + return namelist + + +def extract(src, dest=None): + """ + Takes in a tar or zip file and extracts it to dest + + If dest is not specified, extracts to os.path.dirname(src) + + Returns the list of top level files that were extracted + """ + + assert os.path.exists(src), "'%s' does not exist" % src + + if dest is None: + dest = os.path.dirname(src) + elif not os.path.isdir(dest): + os.makedirs(dest) + assert not os.path.isfile(dest), "dest cannot be a file" + + if zipfile.is_zipfile(src): + namelist = extract_zip(src, dest) + elif tarfile.is_tarfile(src): + namelist = extract_tarball(src, dest) + else: + raise Exception("mozfile.extract: no archive format found for '%s'" % + src) + + # namelist returns paths with forward slashes even in windows + top_level_files = [os.path.join(dest, name.rstrip('/')) for name in namelist + if len(name.rstrip('/').split('/')) == 1] + + # namelist doesn't include folders, append these to the list + for name in namelist: + index = name.find('/') + if index != -1: + root = os.path.join(dest, name[:index]) + if root not in top_level_files: + top_level_files.append(root) + + return top_level_files + + +### utilities for removal of files and directories + +def rmtree(dir): + """Deprecated wrapper method to remove a directory tree. + + Ensure to update your code to use mozfile.remove() directly + + :param dir: directory to be removed + """ + + return remove(dir) + + +def remove(path): + """Removes the specified file, link, or directory tree + + This is a replacement for shutil.rmtree that works better under + windows. + + :param path: path to be removed + """ + + def _call_with_windows_retry(func, path, retry_max=5, retry_delay=0.5): + """ + It's possible to see spurious errors on Windows due to various things + keeping a handle to the directory open (explorer, virus scanners, etc) + So we try a few times if it fails with a known error. + """ + retry_count = 0 + while True: + try: + func(path) + break + except WindowsError as e: + # Error 5 == Access is denied + # Error 32 == The process cannot access the file because it is + # being used by another process + # Error 145 == The directory is not empty + + if retry_count == retry_max or e.winerror not in [5, 32, 145]: + raise + retry_count += 1 + + print 'Retrying to remove "%s" because it is in use.' % path + time.sleep(retry_delay) + + if not os.path.exists(path): + return + + path_stats = os.stat(path) + + if os.path.isfile(path) or os.path.islink(path): + # Verify the file or link is read/write for the current user + os.chmod(path, path_stats.st_mode | stat.S_IRUSR | stat.S_IWUSR) + _call_with_windows_retry(os.remove, path) + + elif os.path.isdir(path): + # Verify the directory is read/write/execute for the current user + os.chmod(path, path_stats.st_mode | stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) + _call_with_windows_retry(shutil.rmtree, path) + +def depth(directory): + """returns the integer depth of a directory or path relative to '/' """ + + directory = os.path.abspath(directory) + level = 0 + while True: + directory, remainder = os.path.split(directory) + level += 1 + if not remainder: + break + return level + +# ASCII delimeters +ascii_delimeters = { + 'vertical_line' : '|', + 'item_marker' : '+', + 'last_child' : '\\' + } + +# unicode delimiters +unicode_delimeters = { + 'vertical_line' : '│', + 'item_marker' : '├', + 'last_child' : '└' + } + +def tree(directory, + item_marker=unicode_delimeters['item_marker'], + vertical_line=unicode_delimeters['vertical_line'], + last_child=unicode_delimeters['last_child'], + sort_key=lambda x: x.lower()): + """ + display tree directory structure for `directory` + """ + + retval = [] + indent = [] + last = {} + top = depth(directory) + + for dirpath, dirnames, filenames in os.walk(directory, topdown=True): + + abspath = os.path.abspath(dirpath) + basename = os.path.basename(abspath) + parent = os.path.dirname(abspath) + level = depth(abspath) - top + + # sort articles of interest + for resource in (dirnames, filenames): + resource[:] = sorted(resource, key=sort_key) + + files_end = item_marker + dirpath_marker = item_marker + + if level > len(indent): + indent.append(vertical_line) + indent = indent[:level] + + if dirnames: + files_end = item_marker + last[abspath] = dirnames[-1] + else: + files_end = last_child + + if last.get(parent) == os.path.basename(abspath): + # last directory of parent + dirpath_mark = last_child + indent[-1] = ' ' + elif not indent: + dirpath_mark = '' + else: + dirpath_mark = item_marker + + # append the directory and piece of tree structure + # if the top-level entry directory, print as passed + retval.append('%s%s%s'% (''.join(indent[:-1]), + dirpath_mark, + basename if retval else directory)) + # add the files + if filenames: + last_file = filenames[-1] + retval.extend([('%s%s%s' % (''.join(indent), + files_end if filename == last_file else item_marker, + filename)) + for index, filename in enumerate(filenames)]) + + return '\n'.join(retval) + + +### utilities for temporary resources + +class NamedTemporaryFile(object): + """ + Like tempfile.NamedTemporaryFile except it works on Windows + in the case where you open the created file a second time. + + This behaves very similarly to tempfile.NamedTemporaryFile but may + not behave exactly the same. For example, this function does not + prevent fd inheritance by children. + + Example usage: + + with NamedTemporaryFile() as fh: + fh.write(b'foobar') + + print('Filename: %s' % fh.name) + + see https://bugzilla.mozilla.org/show_bug.cgi?id=821362 + """ + def __init__(self, mode='w+b', bufsize=-1, suffix='', prefix='tmp', + dir=None, delete=True): + + fd, path = tempfile.mkstemp(suffix, prefix, dir, 't' in mode) + os.close(fd) + + self.file = open(path, mode) + self._path = path + self._delete = delete + self._unlinked = False + + def __getattr__(self, k): + return getattr(self.__dict__['file'], k) + + def __iter__(self): + return self.__dict__['file'] + + def __enter__(self): + self.file.__enter__() + return self + + def __exit__(self, exc, value, tb): + self.file.__exit__(exc, value, tb) + if self.__dict__['_delete']: + os.unlink(self.__dict__['_path']) + self._unlinked = True + + def __del__(self): + if self.__dict__['_unlinked']: + return + self.file.__exit__(None, None, None) + if self.__dict__['_delete']: + os.unlink(self.__dict__['_path']) + + +@contextmanager +def TemporaryDirectory(): + """ + create a temporary directory using tempfile.mkdtemp, and then clean it up. + + Example usage: + with TemporaryDirectory() as tmp: + open(os.path.join(tmp, "a_temp_file"), "w").write("data") + + """ + tempdir = tempfile.mkdtemp() + try: + yield tempdir + finally: + shutil.rmtree(tempdir) + + +### utilities dealing with URLs + +def is_url(thing): + """ + Return True if thing looks like a URL. + """ + + parsed = urlparse.urlparse(thing) + if 'scheme' in parsed: + return len(parsed.scheme) >= 2 + else: + return len(parsed[0]) >= 2 + +def load(resource): + """ + open a file or URL for reading. If the passed resource string is not a URL, + or begins with 'file://', return a ``file``. Otherwise, return the + result of urllib2.urlopen() + """ + + # handle file URLs separately due to python stdlib limitations + if resource.startswith('file://'): + resource = resource[len('file://'):] + + if not is_url(resource): + # if no scheme is given, it is a file path + return file(resource) + + return urllib2.urlopen(resource) + |