wizard/util.py

   1 """
   2 Miscellaneous utility functions and classes.
   3
   4 .. testsetup:: *
   5
   6     from wizard.util import *
   7 """
   8
   9 import os.path
  10 import os
  11 import subprocess
  12 import pwd
  13 import sys
  14
  15 import wizard
  16
  17 class ChangeDirectory(object):
  18     """
  19     Context for temporarily changing the working directory.
  20
  21         >>> with ChangeDirectory("/tmp"):
  22         ...    print os.getcwd()
  23         /tmp
  24     """
  25     def __init__(self, dir):
  26         self.dir = dir
  27         self.olddir = None
  28     def __enter__(self):
  29         self.olddir = os.getcwd()
  30         os.chdir(self.dir)
  31     def __exit__(self, *args):
  32         os.chdir(self.olddir)
  33
  34 class Counter(object):
  35     """
  36     Object for counting different values when you don't know what
  37     they are a priori.  Supports index access and iteration.
  38
  39         >>> counter = Counter()
  40         >>> counter.count("foo")
  41         >>> print counter["foo"]
  42         1
  43     """
  44     def __init__(self):
  45         self.dict = {}
  46     def count(self, value):
  47         """Increments count for ``value``."""
  48         self.dict.setdefault(value, 0)
  49         self.dict[value] += 1
  50     def __getitem__(self, key):
  51         return self.dict[key]
  52     def __iter__(self):
  53         return self.dict.__iter__()
  54
  55 def dictmap(f, d):
  56     """
  57     A map function for dictionaries.  Only changes values.
  58
  59         >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
  60         {'a': 3, 'b': 4}
  61     """
  62     return dict((k,f(v)) for k,v in d.items())
  63
  64 def get_exception_name(output):
  65     """
  66     Reads the traceback from a Python program and grabs the
  67     fully qualified exception name.
  68     """
  69     lines = output.split("\n")
  70     for line in lines[1:]: # skip the "traceback" line
  71         line = line.rstrip()
  72         if line[0] == ' ': continue
  73         if line[-1] == ":":
  74             return line[:-1]
  75         else:
  76             return line
  77
  78 def get_dir_uid(dir):
  79     """Finds the uid of the person who owns this directory."""
  80     return os.stat(dir).st_uid
  81
  82 def get_dir_owner(dir = "."):
  83     """
  84     Finds the name of the locker this directory is in.
  85
  86     .. note::
  87
  88         This function uses the passwd database and thus
  89         only works on scripts servers.
  90     """
  91     return pwd.getpwuid(get_dir_uid(dir)).pw_name
  92
  93 def get_revision():
  94     """Returns the commit ID of the current Wizard install."""
  95     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
  96     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
  97
  98 def get_operator_info():
  99     """
 100     Returns tuple of ``(realname, email)`` from Hesiod about the
 101     this operator of this script from :func:`get_operator_name`.
 102     Useful when generating commit messages.
 103     """
 104     username = get_operator_name()
 105     hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
 106     fields = hesinfo.partition(",")[0]
 107     realname = fields.rpartition(":")[2]
 108     return realname, username + "@mit.edu"
 109
 110 def get_operator_git():
 111     """
 112     Returns ``Real Name <username@mit.edu>`` suitable for use in
 113     Git ``Something-by:`` string.
 114     """
 115     return "%s <%s>" % get_operator_info()
 116
 117 def get_operator_name():
 118     """
 119     Returns username of the person operating this script based
 120     off of the :envvar:`SSH_GSSAPI_NAME` environment variable.  Throws
 121     :exc:`NoOperatorInfo` if environment variable is not available.
 122
 123     .. note::
 124
 125         :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
 126         distributions.  Scripts servers are patched to support this
 127         environment variable.
 128     """
 129     principal = os.getenv("SSH_GSSAPI_NAME")
 130     if not principal: raise NoOperatorInfo
 131     instance, _, _ = principal.partition("@")
 132     if instance.endswith("/root"):
 133         username, _, _ = principal.partition("/")
 134     else:
 135         username = instance
 136     return username
 137
 138 def set_operator_env():
 139     """
 140     Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
 141     environment variables if applicable.  Does nothing if
 142     :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
 143     """
 144     try:
 145         op_realname, op_email = get_operator_info()
 146         os.putenv("GIT_COMMITTER_NAME", op_realname)
 147         os.putenv("GIT_COMMITTER_EMAIL", op_email)
 148     except NoOperatorInfo:
 149         pass
 150
 151 def set_author_env():
 152     """
 153     Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
 154     variables if applicable. Does nothing if :func:`get_dir_owner` fails.
 155     """
 156     try:
 157         lockername = get_dir_owner()
 158         os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
 159         os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
 160     except KeyError: # XXX: This doesn't actually make sense
 161         pass
 162
 163 def set_git_env():
 164     """Sets all appropriate environment variables for Git commits."""
 165     set_operator_env()
 166     set_author_env()
 167
 168 def get_git_footer():
 169     """Returns strings for placing in Git log info about Wizard."""
 170     return "\n".join(["Wizard-revision: %s" % get_revision()
 171         ,"Wizard-args: %s" % " ".join(sys.argv)
 172         ])
 173
 174 class NoOperatorInfo(wizard.Error):
 175     """No information could be found about the operator from Kerberos."""
 176     pass
 177