Revamp 'wizard summary', and start writing tutorial docs.

[wizard.git] / wizard / util.py

diff --git a/wizard/util.py b/wizard/util.py
index 2d4e4ff4d1993268dbd0822fd326c7e42049c504..0cd8fc88cb8447fbd960a0d70cc64bbe3c536964 100644 (file)
--- a/wizard/util.py
+++ b/wizard/util.py
@@ -12,6 +12,11 @@ import subprocess
 import pwd
 import sys
 import socket
+import errno
+import itertools
+import signal
+import httplib
+import urllib
 
 import wizard
 
@@ -28,9 +33,9 @@ class ChangeDirectory(object):
         self.olddir = None
     def __enter__(self):
         self.olddir = os.getcwd()
-        os.chdir(self.dir)
+        chdir(self.dir)
     def __exit__(self, *args):
-        os.chdir(self.olddir)
+        chdir(self.olddir)
 
 class Counter(object):
     """
@@ -52,6 +57,15 @@ class Counter(object):
         return self.dict[key]
     def __iter__(self):
         return self.dict.__iter__()
+    def max(self):
+        """Returns the max counter value seen."""
+        return max(self.dict.values())
+    def sum(self):
+        """Returns the sum of all counter values."""
+        return sum(self.dict.values())
+    def keys(self):
+        """Returns the keys of counters."""
+        return self.dict.keys()
 
 class PipeToLess(object):
     """
@@ -68,6 +82,51 @@ class PipeToLess(object):
             self.proc.wait()
             sys.stdout = self.old_stdout
 
+class IgnoreKeyboardInterrupts(object):
+    """
+    Context for temporarily ignoring keyboard interrupts.  Use this
+    if aborting would cause more harm than finishing the job.
+    """
+    def __enter__(self):
+        signal.signal(signal.SIGINT,signal.SIG_IGN)
+    def __exit__(self, *args):
+        signal.signal(signal.SIGINT, signal.default_int_handler)
+
+class LockDirectory(object):
+    """
+    Context for locking a directory.
+    """
+    def __init__(self, lockfile):
+        self.lockfile = lockfile
+    def __enter__(self):
+        try:
+            os.open(self.lockfile, os.O_CREAT | os.O_EXCL)
+        except OSError as e:
+            if e.errno == errno.EEXIST:
+                raise DirectoryLockedError(os.getcwd())
+            elif e.errno == errno.EACCES:
+                raise PermissionsError(os.getcwd())
+            raise
+    def __exit__(self, *args):
+        try:
+            os.unlink(self.lockfile)
+        except OSError:
+            pass
+
+def chdir(dir):
+    """
+    Changes a directory, but has special exceptions for certain
+    classes of errors.
+    """
+    try:
+        os.chdir(dir)
+    except OSError as e:
+        if e.errno == errno.EACCES:
+            raise PermissionsError()
+        elif e.errno == errno.ENOENT:
+            raise NoSuchDirectoryError()
+        else: raise e
+
 def dictmap(f, d):
     """
     A map function for dictionaries.  Only changes values.
@@ -82,7 +141,7 @@ def dictkmap(f, d):
     A map function for dictionaries that passes key and value.
 
         >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
-        {1: 5, 3: 6}
+        {1: 5, 3: 7}
     """
     return dict((k,f(k,v)) for k,v in d.items())
 
@@ -93,6 +152,7 @@ def get_exception_name(output):
     """
     lines = output.split("\n")
     cue = False
+    result = "(unknown)"
     for line in lines[1:]:
         line = line.rstrip()
         if not line: continue
@@ -101,10 +161,7 @@ def get_exception_name(output):
             continue
         if cue:
             cue = False
-            if line[-1] == ":":
-                result = line[:-1]
-            else:
-                result = line
+            return line.partition(':')[0]
     return result
 
 def get_dir_uid(dir):
@@ -121,9 +178,13 @@ def get_dir_owner(dir = "."):
         only works on scripts servers when querying directories
         that live on AFS.
     """
-    pwentry = pwd.getpwuid(get_dir_uid(dir))
-    # XXX: Error handling!
-    return pwentry.pw_name
+    uid = get_dir_uid(dir)
+    try:
+        pwentry = pwd.getpwuid(uid)
+        return pwentry.pw_name
+    except KeyError:
+        # do an pts query to get the name
+        return subprocess.Popen(['pts', 'examine', str(uid)], stdout=subprocess.PIPE).communicate()[0].partition(",")[0].partition(": ")[2]
 
 def get_revision():
     """Returns the commit ID of the current Wizard install."""
@@ -230,7 +291,49 @@ def get_git_footer():
         ,"Wizard-args: %s" % " ".join(sys.argv)
         ])
 
+def safe_unlink(file):
+    """Moves a file/dir to a backup location."""
+    if not os.path.exists(file):
+        return None
+    prefix = "%s.bak" % file
+    name = None
+    for i in itertools.count():
+        name = "%s.%d" % (prefix, i)
+        if not os.path.exists(name):
+            break
+    os.rename(file, name)
+    return name
+
+def fetch(host, path, subpath, post=None):
+    h = httplib.HTTPConnection(host)
+    fullpath = path.rstrip("/") + "/" + subpath.lstrip("/") # to be lenient about input we accept
+    if post:
+        headers = {"Content-type": "application/x-www-form-urlencoded"}
+        h.request("POST", fullpath, urllib.urlencode(post), headers)
+    else:
+        h.request("GET", fullpath)
+    r = h.getresponse()
+    data = r.read()
+    h.close()
+    return data
+
 class NoOperatorInfo(wizard.Error):
     """No information could be found about the operator from Kerberos."""
     pass
 
+class PermissionsError(IOError):
+    errno = errno.EACCES
+
+class NoSuchDirectoryError(IOError):
+    errno = errno.ENOENT
+
+class DirectoryLockedError(wizard.Error):
+    def __init__(self, dir):
+        self.dir = dir
+    def __str__(self):
+        return """
+
+ERROR: Could not acquire lock on directory.  Maybe there is
+another migration process running?
+"""
+