2 Miscellaneous utility functions and classes.
6 from wizard.util import *
23 class ChangeDirectory(object):
25 Context for temporarily changing the working directory.
27 >>> with ChangeDirectory("/tmp"):
31 def __init__(self, dir):
35 self.olddir = os.getcwd()
37 def __exit__(self, *args):
40 class Counter(object):
42 Object for counting different values when you don't know what
43 they are a priori. Supports index access and iteration.
45 >>> counter = Counter()
46 >>> counter.count("foo")
47 >>> print counter["foo"]
52 def count(self, value):
53 """Increments count for ``value``."""
54 self.dict.setdefault(value, 0)
56 def __getitem__(self, key):
59 return self.dict.__iter__()
61 class PipeToLess(object):
63 Context for printing output to a pager. Use this if output
64 is expected to be long.
67 self.proc = subprocess.Popen("less", stdin=subprocess.PIPE)
68 self.old_stdout = sys.stdout
69 sys.stdout = self.proc.stdin
70 def __exit__(self, *args):
72 self.proc.stdin.close()
74 sys.stdout = self.old_stdout
76 class IgnoreKeyboardInterrupts(object):
78 Context for temporarily ignoring keyboard interrupts. Use this
79 if aborting would cause more harm than finishing the job.
82 signal.signal(signal.SIGINT,signal.SIG_IGN)
83 def __exit__(self, *args):
84 signal.signal(signal.SIGINT, signal.default_int_handler)
86 class LockDirectory(object):
88 Context for locking a directory.
90 def __init__(self, lockfile):
91 self.lockfile = lockfile
94 os.open(self.lockfile, os.O_CREAT | os.O_EXCL)
96 if e.errno == errno.EEXIST:
97 raise DirectoryLockedError(os.getcwd())
98 elif e.errno == errno.EACCES:
99 raise command.PermissionsError(os.getcwd())
101 def __exit__(self, *args):
103 os.unlink(self.lockfile)
109 Changes a directory, but has special exceptions for certain
115 if e.errno == errno.EACCES:
116 raise PermissionsError()
117 elif e.errno == errno.ENOENT:
118 raise NoSuchDirectoryError()
123 A map function for dictionaries. Only changes values.
125 >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
128 return dict((k,f(v)) for k,v in d.items())
132 A map function for dictionaries that passes key and value.
134 >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
137 return dict((k,f(k,v)) for k,v in d.items())
139 def get_exception_name(output):
141 Reads the traceback from a Python program and grabs the
142 fully qualified exception name.
144 lines = output.split("\n")
147 for line in lines[1:]:
149 if not line: continue
162 def get_dir_uid(dir):
163 """Finds the uid of the person who owns this directory."""
164 return os.stat(dir).st_uid
166 def get_dir_owner(dir = "."):
168 Finds the name of the locker this directory is in.
172 This function uses the passwd database and thus
173 only works on scripts servers when querying directories
176 pwentry = pwd.getpwuid(get_dir_uid(dir))
177 # XXX: Error handling!
178 return pwentry.pw_name
181 """Returns the commit ID of the current Wizard install."""
182 # If you decide to convert this to use wizard.shell, be warned
183 # that there is a circular dependency, so this function would
184 # probably have to live somewhere else, probably wizard.git
185 wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
186 return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
188 def get_operator_info():
190 Returns tuple of ``(realname, email)`` about the person running
191 the script. If run from a scripts server, get info from Hesiod.
192 Otherwise, use the passwd database (email generated probably won't
193 actually accept mail). Useful when generating commit messages.
195 username = get_operator_name_from_gssapi()
198 hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
199 fields = hesinfo.partition(",")[0]
200 realname = fields.rpartition(":")[2]
201 return realname, username + "@mit.edu"
203 # more traditional approach, but the email probably doesn't work
206 # since root isn't actually a useful designation, but maybe
207 # SUDO_USER contains something helpful
208 sudo_user = os.getenv("SUDO_USER")
211 pwdentry = pwd.getpwnam(sudo_user)
213 pwdentry = pwd.getpwuid(uid)
214 # XXX: error checking might be nice
215 # We follow the Ubuntu convention of gecos being a comma split field
216 # with the person's realname being the first entry.
217 return pwdentry.pw_gecos.split(",")[0], pwdentry.pw_name + "@" + socket.gethostname()
219 def get_operator_git():
221 Returns ``Real Name <username@mit.edu>`` suitable for use in
222 Git ``Something-by:`` string.
224 return "%s <%s>" % get_operator_info()
226 def get_operator_name_from_gssapi():
228 Returns username of the person operating this script based
229 off of the :envvar:`SSH_GSSAPI_NAME` environment variable.
233 :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
234 distributions. Scripts servers are patched to support this
235 environment variable.
237 principal = os.getenv("SSH_GSSAPI_NAME")
240 instance, _, _ = principal.partition("@")
241 if instance.endswith("/root"):
242 username, _, _ = principal.partition("/")
247 def set_operator_env():
249 Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
250 environment variables if applicable. Does nothing if
251 :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
254 op_realname, op_email = get_operator_info()
255 os.putenv("GIT_COMMITTER_NAME", op_realname)
256 os.putenv("GIT_COMMITTER_EMAIL", op_email)
257 except NoOperatorInfo:
260 def set_author_env():
262 Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
263 variables if applicable. Does nothing if :func:`get_dir_owner` fails.
266 # XXX: should check if the directory is in AFS, and if not, use
267 # a more traditional metric
268 lockername = get_dir_owner()
269 os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
270 os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
271 except KeyError: # XXX: This doesn't actually make sense
275 """Sets all appropriate environment variables for Git commits."""
279 def get_git_footer():
280 """Returns strings for placing in Git log info about Wizard."""
281 return "\n".join(["Wizard-revision: %s" % get_revision()
282 ,"Wizard-args: %s" % " ".join(sys.argv)
285 def safe_unlink(file):
286 """Moves a file/dir to a backup location."""
287 if not os.path.exists(file):
289 prefix = "%s.bak" % file
291 for i in itertools.count():
292 name = "%s.%d" % (prefix, i)
293 if not os.path.exists(name):
295 os.rename(file, name)
298 def fetch(host, path, subpath, post=None):
299 h = httplib.HTTPConnection(host)
300 fullpath = path + "/" + subpath
302 headers = {"Content-type": "application/x-www-form-urlencoded"}
303 h.request("POST", fullpath, urllib.urlencode(post), headers)
305 h.request("GET", fullpath)
311 class NoOperatorInfo(wizard.Error):
312 """No information could be found about the operator from Kerberos."""
315 class PermissionsError(IOError):
318 class NoSuchDirectoryError(IOError):
321 class DirectoryLockedError(wizard.Error):
322 def __init__(self, dir):
327 ERROR: Could not acquire lock on directory. Maybe there is
328 another migration process running?