codekingpro/portable-devtools

1import contextlib2import os3import sys4import random5import string6 7 8class ImportKiller:9    "Context manager to make an import of a given name or names fail."10 11    def __init__(self, *names):12        self.names = names13 14    def find_module(self, fullname, path=None):15        if fullname in self.names:16            return self17 18    def load_module(self, fullname):19        assert fullname in self.names20        raise ImportError(fullname)21 22    def __enter__(self):23        self.original = {}24        for name in self.names:25            self.original[name] = sys.modules.pop(name, None)26        sys.meta_path.insert(0, self)27 28    def __exit__(self, *args):29        sys.meta_path.remove(self)30        for key, value in self.original.items():31            if value is not None:32                sys.modules[key] = value33 34 35@contextlib.contextmanager36def NoNoneDictMutator(destination, **changes):37    """Helper context manager to make and unmake changes to a dict.38 39    A None is not a valid value for the destination, and so means that the40    associated name should be removed."""41    original = {}42    for key, value in changes.items():43        original[key] = destination.get(key)44        if value is None:45            if key in destination:46                del destination[key]47        else:48            destination[key] = value49    yield50    for key, value in original.items():51        if value is None:52            if key in destination:53                del destination[key]54        else:55            destination[key] = value56 57 58def Environ(**changes):59    """A context manager to temporarily change the os.environ"""60    return NoNoneDictMutator(os.environ, **changes)61 62 63ALPHABET = string.ascii_letters + string.digits64 65 66def random_string(k, source=ALPHABET):67    """Generate a random string with length <i>k</i>"""68    result = ''69    for i in range(0, k):70        result += random.choice(source)71    return result72