]> scripts.mit.edu Git - wizard.git/blobdiff - wizard/util.py
Implement 'wizard install', with other improvements.
[wizard.git] / wizard / util.py
index 6303219dbff90812f4f3943cf3e3fd9423748148..2b244de5cbbd4c3194ecf68a463110accea46673 100644 (file)
@@ -1,3 +1,11 @@
+"""
+Miscellaneous utility functions and classes.
+
+.. testsetup:: *
+
+    from wizard.util import *
+"""
+
 import os.path
 import os
 import subprocess
@@ -7,7 +15,13 @@ import sys
 import wizard
 
 class ChangeDirectory(object):
-    """Context for temporarily changing directory"""
+    """
+    Context for temporarily changing the working directory.
+
+        >>> with ChangeDirectory("/tmp"):
+        ...    print os.getcwd()
+        /tmp
+    """
     def __init__(self, dir):
         self.dir = dir
         self.olddir = None
@@ -17,9 +31,41 @@ class ChangeDirectory(object):
     def __exit__(self, *args):
         os.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 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 get_exception_name(output):
-    """Reads the stderr output of another Python command and grabs the
-    fully qualified exception name"""
+    """
+    Reads the traceback from a Python program and grabs the
+    fully qualified exception name.
+    """
     lines = output.split("\n")
     for line in lines[1:]: # skip the "traceback" line
         line = line.rstrip()
@@ -33,8 +79,15 @@ def get_dir_uid(dir):
     """Finds the uid of the person who owns this directory."""
     return os.stat(dir).st_uid
 
-def get_dir_owner(dir):
-    """Finds the name of the locker this directory is in."""
+def get_dir_owner(dir = "."):
+    """
+    Finds the name of the locker this directory is in.
+
+    .. note::
+
+        This function uses the passwd database and thus
+        only works on scripts servers.
+    """
     return pwd.getpwuid(get_dir_uid(dir)).pw_name
 
 def get_revision():
@@ -43,9 +96,11 @@ def get_revision():
     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
 
 def get_operator_info():
-    """Returns tuple of (realname, email) of person who is operating
-    this script, as told to use by the Kerberos principal name.
-    Useful for commit messages."""
+    """
+    Returns tuple of ``(realname, email)`` from Hesiod about the
+    this operator of this script from :func:`get_operator_name`.
+    Useful when generating commit messages.
+    """
     username = get_operator_name()
     hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
     fields = hesinfo.partition(",")[0]
@@ -53,14 +108,26 @@ def get_operator_info():
     return realname, username + "@mit.edu"
 
 def get_operator_git():
-    """Returns Real Name <username@mit.edu> suitable for use in
-    Git Something-by: string."""
+    """
+    Returns ``Real Name <username@mit.edu>`` suitable for use in
+    Git ``Something-by:`` string.
+    """
     return "%s <%s>" % get_operator_info()
 
 def get_operator_name():
-    """Returns username of the person operating this script."""
+    """
+    Returns username of the person operating this script based
+    off of the :envvar:`SSH_GSSAPI_NAME` environment variable.  Throws
+    :exc:`NoOperatorInfo` if environment variable is not available.
+
+    .. note::
+
+        :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
+        distributions.  Scripts servers are patched to support this
+        environment variable.
+    """
     principal = os.getenv("SSH_GSSAPI_NAME")
-    if not principal: raise NoOperatorInfo()
+    if not principal: raise NoOperatorInfo
     instance, _, _ = principal.partition("@")
     if instance.endswith("/root"):
         username, _, _ = principal.partition("/")
@@ -69,6 +136,11 @@ def get_operator_name():
     return username
 
 def set_operator_env():
+    """
+    Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
+    environment variables if applicable.  Does nothing if
+    :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
+    """
     try:
         op_realname, op_email = get_operator_info()
         os.putenv("GIT_COMMITTER_NAME", op_realname)
@@ -77,25 +149,29 @@ def set_operator_env():
         pass
 
 def set_author_env():
+    """
+    Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
+    variables if applicable. Does nothing if :func:`get_dir_owner` fails.
+    """
     try:
-        lockername = get_dir_owner(".")
+        lockername = get_dir_owner()
         os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
         os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
-    except KeyError:
+    except KeyError: # XXX: This doesn't actually make sense
         pass
 
 def set_git_env():
+    """Sets all appropriate environment variables for Git commits."""
     set_operator_env()
     set_author_env()
 
 def get_git_footer():
+    """Returns strings for placing in Git log info about Wizard."""
     return "\n".join(["Wizard-revision: %s" % get_revision()
         ,"Wizard-args: %s" % " ".join(sys.argv)
         ])
 
-class Error(wizard.Error):
-    pass
-
-class NoOperatorInfo(Error):
+class NoOperatorInfo(wizard.Error):
+    """No information could be found about the operator from Kerberos."""
     pass