Improve upgrade dry run, improve error messages.

[wizard.git] / wizard / util.py

diff --git a/wizard/util.py b/wizard/util.py
index e0694574cd1b7eef47ae3f04a62735f837f0e609..2588d173bba79bd185ee82f18ef62fbeb95b2a1d 100644 (file)
--- a/wizard/util.py
+++ b/wizard/util.py
@@ -2,9 +2,37 @@ import os.path
 import os
 import subprocess
 import pwd
+import sys
 
 import wizard
 
+class ChangeDirectory(object):
+    """Context for temporarily changing directory"""
+    def __init__(self, dir):
+        self.dir = dir
+        self.olddir = None
+    def __enter__(self):
+        self.olddir = os.getcwd()
+        os.chdir(self.dir)
+    def __exit__(self, *args):
+        os.chdir(self.olddir)
+
+class Counter(object):
+    def __init__(self):
+        self.dict = {}
+    def count(self, 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.  Does not allow changing keys, only
+    values"""
+    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"""
@@ -21,11 +49,11 @@ 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):
+def get_dir_owner(dir = "."):
     """Finds the name of the locker this directory is in."""
     return pwd.getpwuid(get_dir_uid(dir)).pw_name
 
-def get_version():
+def get_revision():
     """Returns the commit ID of the current Wizard install."""
     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
@@ -56,6 +84,31 @@ def get_operator_name():
         username = instance
     return username
 
+def set_operator_env():
+    try:
+        op_realname, op_email = get_operator_info()
+        os.putenv("GIT_COMMITTER_NAME", op_realname)
+        os.putenv("GIT_COMMITTER_EMAIL", op_email)
+    except NoOperatorInfo:
+        pass
+
+def set_author_env():
+    try:
+        lockername = get_dir_owner()
+        os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
+        os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
+    except KeyError:
+        pass
+
+def set_git_env():
+    set_operator_env()
+    set_author_env()
+
+def get_git_footer():
+    return "\n".join(["Wizard-revision: %s" % get_revision()
+        ,"Wizard-args: %s" % " ".join(sys.argv)
+        ])
+
 class Error(wizard.Error):
     pass