Added TempDir

expect_tests/tempdir.py

+"""Provides an automatically-cleaned temporary directory.
+
+All functions in this module are thread-safe.
+"""
+
+import glob
+import os
+import shutil
+import tempfile
+
+TEMP_PID_FILE = '.expect_tests_pidfile'
+TEMP_SUFFIX = '.expect_tests_temp'
+
+
+class TempDir(object):
+  """Context manager for a temporary directory.
+
+  Use tempfile.TemporaryDirectory instead with Python 3.
+  """
+  def __init__(self):
+    self.temp_dir = None
+
+  def __enter__(self):
+    self.temp_dir = set_temp_dir()
+    return self.temp_dir
+
+  def __exit__(self, _type, _value, _traceback):
+    clear_temp_dir(self.temp_dir)
+
+
+def set_temp_dir(suffix=TEMP_SUFFIX):
+  """Provides a temporary directory for tests.
+
+  This function also takes care of removing old temp directories from dead
+  processes.
+  """
+  temp_dir = tempfile.mkdtemp(suffix)
+
+  # Clean up obsolete directories.
+  for p in glob.glob(os.path.join(os.path.dirname(temp_dir), '*'+suffix)):
+    if p == temp_dir:
+      continue
+
+    pfile = os.path.join(p, TEMP_PID_FILE)
+    if os.path.exists(pfile):
+      with open(pfile, 'rb') as f:
+        try:
+          os.kill(int(f.read()), 0)
+          continue
+        except (OSError, TypeError):
+          pass
+
+    shutil.rmtree(p, ignore_errors=True)
+
+  # Create current PID file.
+  with open(os.path.join(temp_dir, TEMP_PID_FILE), 'wb') as f:
+    f.write(str(os.getpid()))
+  return temp_dir
+
+
+def clear_temp_dir(temp_dir):
+  # Try to nuke the pidfile first.
+  pfile = os.path.join(temp_dir, TEMP_PID_FILE)
+  try:
+    os.unlink(pfile)
+  except OSError as e:
+    print >> sys.stderr, "Error removing %r: %s" % (pfile, e)
+  shutil.rmtree(temp_dir, ignore_errors=True)