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>

1 files changed, 19 insertions, 4 deletions

diff --git a/mediagoblin/workbench.py b/mediagoblin/workbench.py
index 2331b551..bf18b6cd 100644
--- 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):