+from wizard import user
+
+def boolish(val):
+ """
+ Parse the contents of an environment variable as a boolean.
+ This recognizes more values as ``False`` than :func:`bool` would.
+
+ >>> boolish("0")
+ False
+ >>> boolish("no")
+ False
+ >>> boolish("1")
+ True
+ """
+ try:
+ return bool(int(val))
+ except (ValueError, TypeError):
+ if val == "No" or val == "no" or val == "false" or val == "False":
+ return False
+ return bool(val)
+
+class ChangeDirectory(object):
+ """
+ Context for temporarily changing the working directory.
+
+ >>> with ChangeDirectory("/tmp"):
+ ... print os.getcwd()
+ /tmp
+ """
+ def __init__(self, dir):
+ self.dir = dir
+ self.olddir = None
+ def __enter__(self):
+ self.olddir = os.getcwd()
+ chdir(self.dir)
+ def __exit__(self, *args):
+ chdir(self.olddir)
+
+class Counter(object):
+ """
+ Object for counting different values when you don't know what
+ they are a priori. Supports index access and iteration.
+
+ >>> counter = Counter()
+ >>> counter.count("foo")
+ >>> print counter["foo"]
+ 1
+ """
+ def __init__(self):
+ self.dict = {}
+ def count(self, value):
+ """Increments count for ``value``."""
+ self.dict.setdefault(value, 0)
+ self.dict[value] += 1
+ def __getitem__(self, key):
+ return self.dict[key]
+ def __iter__(self):
+ return self.dict.__iter__()
+ def max(self):
+ """Returns the max counter value seen."""
+ return max(self.dict.values())
+ def sum(self):
+ """Returns the sum of all counter values."""
+ return sum(self.dict.values())
+ def keys(self):
+ """Returns the keys of counters."""
+ return self.dict.keys()
+
+class PipeToLess(object):
+ """
+ Context for printing output to a pager. Use this if output
+ is expected to be long.
+ """
+ def __enter__(self):
+ self.proc = subprocess.Popen("less", stdin=subprocess.PIPE)
+ self.old_stdout = sys.stdout
+ sys.stdout = self.proc.stdin
+ def __exit__(self, *args):
+ if self.proc:
+ self.proc.stdin.close()
+ self.proc.wait()
+ sys.stdout = self.old_stdout
+
+class IgnoreKeyboardInterrupts(object):
+ """
+ Context for temporarily ignoring keyboard interrupts. Use this
+ if aborting would cause more harm than finishing the job.
+ """
+ def __enter__(self):
+ signal.signal(signal.SIGINT,signal.SIG_IGN)
+ def __exit__(self, *args):
+ signal.signal(signal.SIGINT, signal.default_int_handler)
+
+class LockDirectory(object):
+ """
+ Context for locking a directory.
+ """
+ def __init__(self, lockfile, expiry = 3600):
+ self.lockfile = lockfile
+ self.expiry = expiry # by default an hour
+ def __enter__(self):
+ # It's A WAVY
+ for i in range(0, 3):
+ try:
+ os.open(self.lockfile, os.O_CREAT | os.O_EXCL)
+ open(self.lockfile, "w").write("%d" % os.getpid())
+ except OSError as e:
+ if e.errno == errno.EEXIST:
+ # There is a possibility of infinite recursion, but we
+ # expect it to be unlikely, and not harmful if it does happen
+ with LockDirectory(self.lockfile + "_"):
+ # See if we can break the lock
+ try:
+ pid = open(self.lockfile, "r").read().strip()
+ if not os.path.exists("/proc/%s" % pid):
+ # break the lock, try again
+ logging.warning("Breaking orphaned lock at %s", self.lockfile)
+ os.unlink(self.lockfile)
+ continue
+ try:
+ # check if the file is expiry old, if so, break the lock, try again
+ if time.time() - os.stat(self.lockfile).st_mtime > self.expiry:
+ logging.warning("Breaking stale lock at %s", self.lockfile)
+ os.unlink(self.lockfile)
+ continue
+ except OSError as e:
+ if e.errno == errno.ENOENT:
+ continue
+ raise
+ except IOError:
+ # oh hey, it went away; try again
+ continue
+ raise DirectoryLockedError(os.getcwd())
+ elif e.errno == errno.EACCES:
+ raise PermissionsError(os.getcwd())
+ raise
+ return
+ raise DirectoryLockedError(os.getcwd())
+ def __exit__(self, *args):
+ try:
+ os.unlink(self.lockfile)
+ except OSError:
+ pass
+
+def chdir(dir):
+ """
+ Changes a directory, but has special exceptions for certain
+ classes of errors.
+ """
+ try:
+ os.chdir(dir)
+ except OSError as e:
+ if e.errno == errno.EACCES:
+ raise PermissionsError()
+ elif e.errno == errno.ENOENT:
+ raise NoSuchDirectoryError()
+ else: raise e
+
+def dictmap(f, d):
+ """
+ A map function for dictionaries. Only changes values.
+
+ >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
+ {'a': 3, 'b': 4}
+ """
+ return dict((k,f(v)) for k,v in d.items())
+
+def dictkmap(f, d):
+ """
+ A map function for dictionaries that passes key and value.
+
+ >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
+ {1: 5, 3: 7}
+ """
+ return dict((k,f(k,v)) for k,v in d.items())