wizard/util.py

   1 import os.path
   2 import os
   3 import subprocess
   4 import pwd
   5 import sys
   6
   7 import wizard
   8
   9 class ChangeDirectory(object):
  10     """Context for temporarily changing directory"""
  11     def __init__(self, dir):
  12         self.dir = dir
  13         self.olddir = None
  14     def __enter__(self):
  15         self.olddir = os.getcwd()
  16         os.chdir(self.dir)
  17     def __exit__(self, *args):
  18         os.chdir(self.olddir)
  19
  20 class Counter(object):
  21     def __init__(self):
  22         self.dict = {}
  23     def count(self, value):
  24         self.dict.setdefault(value, 0)
  25         self.dict[value] += 1
  26     def __getitem__(self, key):
  27         return self.dict[key]
  28     def __iter__(self):
  29         return self.dict.__iter__()
  30
  31 def dictmap(f, d):
  32     """A map function for dictionaries.  Does not allow changing keys, only
  33     values"""
  34     return dict((k,f(v)) for k,v in d.items())
  35
  36 def get_exception_name(output):
  37     """Reads the stderr output of another Python command and grabs the
  38     fully qualified exception name"""
  39     lines = output.split("\n")
  40     for line in lines[1:]: # skip the "traceback" line
  41         line = line.rstrip()
  42         if line[0] == ' ': continue
  43         if line[-1] == ":":
  44             return line[:-1]
  45         else:
  46             return line
  47
  48 def get_dir_uid(dir):
  49     """Finds the uid of the person who owns this directory."""
  50     return os.stat(dir).st_uid
  51
  52 def get_dir_owner(dir = "."):
  53     """Finds the name of the locker this directory is in."""
  54     return pwd.getpwuid(get_dir_uid(dir)).pw_name
  55
  56 def get_revision():
  57     """Returns the commit ID of the current Wizard install."""
  58     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
  59     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
  60
  61 def get_operator_info():
  62     """Returns tuple of (realname, email) of person who is operating
  63     this script, as told to use by the Kerberos principal name.
  64     Useful for commit messages."""
  65     username = get_operator_name()
  66     hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
  67     fields = hesinfo.partition(",")[0]
  68     realname = fields.rpartition(":")[2]
  69     return realname, username + "@mit.edu"
  70
  71 def get_operator_git():
  72     """Returns Real Name <username@mit.edu> suitable for use in
  73     Git Something-by: string."""
  74     return "%s <%s>" % get_operator_info()
  75
  76 def get_operator_name():
  77     """Returns username of the person operating this script."""
  78     principal = os.getenv("SSH_GSSAPI_NAME")
  79     if not principal: raise NoOperatorInfo()
  80     instance, _, _ = principal.partition("@")
  81     if instance.endswith("/root"):
  82         username, _, _ = principal.partition("/")
  83     else:
  84         username = instance
  85     return username
  86
  87 def set_operator_env():
  88     try:
  89         op_realname, op_email = get_operator_info()
  90         os.putenv("GIT_COMMITTER_NAME", op_realname)
  91         os.putenv("GIT_COMMITTER_EMAIL", op_email)
  92     except NoOperatorInfo:
  93         pass
  94
  95 def set_author_env():
  96     try:
  97         lockername = get_dir_owner()
  98         os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
  99         os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
 100     except KeyError:
 101         pass
 102
 103 def set_git_env():
 104     set_operator_env()
 105     set_author_env()
 106
 107 def get_git_footer():
 108     return "\n".join(["Wizard-revision: %s" % get_revision()
 109         ,"Wizard-args: %s" % " ".join(sys.argv)
 110         ])
 111
 112 class Error(wizard.Error):
 113     pass
 114
 115 class NoOperatorInfo(Error):
 116     pass
 117