Make Workbench() a context manager

This allows us to use "with Workbench() as foo: do_stuff..."
No consumers have been switched yet though.

Signed-off-by: Sebastian Spaeth <Sebastian@SSpaeth.de>

diff --git a/mediagoblin/workbench.py b/mediagoblin/workbench.py
index 2331b5511e59ffa377e856ab0a14b0cd8649167c..bf18b6cdf9032bc582c2ad505152e97e2054cb39 100644 (file)
--- a/mediagoblin/workbench.py
+++ b/mediagoblin/workbench.py
@@ -127,18 +127,33 @@ class Workbench(object):
         """
         # just in case
         workbench = os.path.abspath(self.dir)
-
         shutil.rmtree(workbench)
-
         del self.dir
 
+    def __enter__(self):
+        """Make Workbench a context manager so we can use `with Workbench() as bench:`"""
+        return self
+
+    def  __exit__(self, *args):
+        """Clean up context manager, aka ourselves, deleting the workbench"""
+        self.destroy_self()
+
 
 class WorkbenchManager(object):
     """
     A system for generating and destroying workbenches.
 
-    Workbenches are actually just subdirectories of a temporary storage space
-    for during the processing stage.
+    Workbenches are actually just subdirectories of a (local) temporary
+    storage space for during the processing stage. The preferred way to
+    create them is to use:
+
+        with workbenchmger.create_workbench as workbench:
+            do stuff...
+
+    This will automatically clean up all temporary directories even in
+    case of an exceptions. Also check the
+    @mediagoblin.decorators.get_workbench decorator for a convenient
+    wrapper.
     """
 
     def __init__(self, base_workbench_dir):