Merge remote branch 'remotes/elrond/dev/workbench_class'

3 files changed, 78 insertions, 67 deletions

diff --git a/mediagoblin/process_media/__init__.py b/mediagoblin/process_media/__init__.py
index a2d843e9..0dce1418 100644
--- a/mediagoblin/process_media/__init__.py
+++ b/mediagoblin/process_media/__init__.py
@@ -39,8 +39,8 @@ def process_media_initial(media_id):
         {'_id': ObjectId(media_id)})
 
     queued_filepath = entry['queued_media_file']
-    queued_filename = mgg.workbench_manager.localized_file(
-        workbench, mgg.queue_store, queued_filepath,
+    queued_filename = workbench.localized_file(
+        mgg.queue_store, queued_filepath,
         'source')
 
     queued_file = file(queued_filename, 'r')
@@ -77,4 +77,4 @@ def process_media_initial(media_id):
     entry.save()
 
     # clean up workbench
-    mgg.workbench_manager.destroy_workbench(workbench)
+    workbench.destroy_self()
diff --git a/mediagoblin/tests/test_workbench.py b/mediagoblin/tests/test_workbench.py
index 89f2ef33..953835a1 100644
--- a/mediagoblin/tests/test_workbench.py
+++ b/mediagoblin/tests/test_workbench.py
@@ -30,29 +30,29 @@ class TestWorkbench(object):
 
     def test_create_workbench(self):
         workbench = self.workbench_manager.create_workbench()
-        assert os.path.isdir(workbench)
-        assert workbench.startswith(self.workbench_manager.base_workbench_dir)
+        assert os.path.isdir(workbench.dir)
+        assert workbench.dir.startswith(self.workbench_manager.base_workbench_dir)
+
+    def test_joinpath(self):
+        this_workbench = self.workbench_manager.create_workbench()
+        tmpname = this_workbench.joinpath('temp.txt')
+        assert tmpname == os.path.join(this_workbench.dir, 'temp.txt')
+        this_workbench.destroy_self()
 
     def test_destroy_workbench(self):
         # kill a workbench
         this_workbench = self.workbench_manager.create_workbench()
-        tmpfile = file(os.path.join(this_workbench, 'temp.txt'), 'w')
+        tmpfile_name = this_workbench.joinpath('temp.txt')
+        tmpfile = file(tmpfile_name, 'w')
         with tmpfile:
             tmpfile.write('lollerskates')
 
-        assert os.path.exists(os.path.join(this_workbench, 'temp.txt'))
-
-        self.workbench_manager.destroy_workbench(this_workbench)
-        assert not os.path.exists(os.path.join(this_workbench, 'temp.txt'))
-        assert not os.path.exists(this_workbench)
-
-        # make sure we can't kill other stuff though
-        dont_kill_this = tempfile.mkdtemp()
+        assert os.path.exists(tmpfile_name)
 
-        assert_raises(
-            workbench.WorkbenchOutsideScope,
-            self.workbench_manager.destroy_workbench,
-            dont_kill_this)
+        wb_dir = this_workbench.dir
+        this_workbench.destroy_self()
+        assert not os.path.exists(tmpfile_name)
+        assert not os.path.exists(wb_dir)
 
     def test_localized_file(self):
         tmpdir, this_storage = get_tmp_filestorage()
@@ -65,8 +65,7 @@ class TestWorkbench(object):
             our_file.write('Our file')
 
         # with a local file storage
-        filename = self.workbench_manager.localized_file(
-            this_workbench, this_storage, filepath)
+        filename = this_workbench.localized_file(this_storage, filepath)
         assert filename == os.path.join(
             tmpdir, 'dir1/dir2/ourfile.txt')
 
@@ -77,20 +76,19 @@ class TestWorkbench(object):
         with this_storage.get_file(filepath, 'w') as our_file:
             our_file.write('Our file')
 
-        filename = self.workbench_manager.localized_file(
-            this_workbench, this_storage, filepath)
+        filename = this_workbench.localized_file(this_storage, filepath)
         assert filename == os.path.join(
-            this_workbench, 'ourfile.txt')
+            this_workbench.dir, 'ourfile.txt')
         
         # fake remote file storage, filename_if_copying set
-        filename = self.workbench_manager.localized_file(
-            this_workbench, this_storage, filepath, 'thisfile')
+        filename = this_workbench.localized_file(
+            this_storage, filepath, 'thisfile')
         assert filename == os.path.join(
-            this_workbench, 'thisfile.txt')
+            this_workbench.dir, 'thisfile.txt')
 
         # fake remote file storage, filename_if_copying set,
         # keep_extension_if_copying set to false
-        filename = self.workbench_manager.localized_file(
-            this_workbench, this_storage, filepath, 'thisfile.text', False)
+        filename = this_workbench.localized_file(
+            this_storage, filepath, 'thisfile.text', False)
         assert filename == os.path.join(
-            this_workbench, 'thisfile.text')
+            this_workbench.dir, 'thisfile.text')
diff --git a/mediagoblin/workbench.py b/mediagoblin/workbench.py
index d7252623..f83c4fa0 100644
--- a/mediagoblin/workbench.py
+++ b/mediagoblin/workbench.py
@@ -23,54 +23,34 @@ DEFAULT_WORKBENCH_DIR = os.path.join(
     tempfile.gettempdir(), u'mgoblin_workbench')
 
 
-# Exception(s)
-# ------------
-
-class WorkbenchOutsideScope(Exception):
-    """
-    Raised when a workbench is outside a WorkbenchManager scope.
-    """
-    pass
-
-
 # Actual workbench stuff
 # ----------------------
 
-class WorkbenchManager(object):
+class Workbench(object):
     """
-    A system for generating and destroying workbenches.
+    Represent the directory for the workbench
 
-    Workbenches are actually just subdirectories of a temporary storage space
-    for during the processing stage.
+    WARNING: DO NOT create Workbench objects on your own,
+    let the WorkbenchManager do that for you!
     """
-
-    def __init__(self, base_workbench_dir):
-        self.base_workbench_dir = os.path.abspath(base_workbench_dir)
-        if not os.path.exists(self.base_workbench_dir):
-            os.makedirs(self.base_workbench_dir)
-        
-    def create_workbench(self):
+    def __init__(self, dir):
         """
-        Create and return the path to a new workbench (directory).
+        WARNING: DO NOT create Workbench objects on your own,
+        let the WorkbenchManager do that for you!
         """
-        return tempfile.mkdtemp(dir=self.base_workbench_dir)
+        self.dir = dir
 
-    def destroy_workbench(self, workbench):
-        """
-        Destroy this workbench!  Deletes the directory and all its contents!
+    def __unicode__(self):
+        return unicode(self.dir)
+    def __str__(self):
+        return str(self.dir)
+    def __repr__(self):
+        return repr(self.dir)
 
-        Makes sure the workbench actually belongs to this manager though.
-        """
-        # just in case
-        workbench = os.path.abspath(workbench)
-
-        if not workbench.startswith(self.base_workbench_dir):
-            raise WorkbenchOutsideScope(
-                "Can't destroy workbench outside the base workbench dir")
-
-        shutil.rmtree(workbench)
+    def joinpath(self, *args):
+        return os.path.join(self.dir, *args)
 
-    def localized_file(self, workbench, storage, filepath,
+    def localized_file(self, storage, filepath,
                        filename_if_copying=None,
                        keep_extension_if_copying=True):
         """
@@ -126,10 +106,43 @@ class WorkbenchManager(object):
                     dest_filename = filename_if_copying
 
             full_dest_filename = os.path.join(
-                workbench, dest_filename)
+                self.dir, dest_filename)
 
             # copy it over
             storage.copy_locally(
                 filepath, full_dest_filename)
 
             return full_dest_filename
+
+    def destroy_self(self):
+        """
+        Destroy this workbench!  Deletes the directory and all its contents!
+
+        WARNING: Does no checks for a sane value in self.dir!
+        """
+        # just in case
+        workbench = os.path.abspath(self.dir)
+
+        shutil.rmtree(workbench)
+
+        del self.dir
+
+
+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.
+    """
+
+    def __init__(self, base_workbench_dir):
+        self.base_workbench_dir = os.path.abspath(base_workbench_dir)
+        if not os.path.exists(self.base_workbench_dir):
+            os.makedirs(self.base_workbench_dir)
+        
+    def create_workbench(self):
+        """
+        Create and return the path to a new workbench (directory).
+        """
+        return Workbench(tempfile.mkdtemp(dir=self.base_workbench_dir))