Fixed a bug that had binary stls be handled by the ascii stl parser.
[mediagoblin.git] / mediagoblin / media_types / stl / processing.py
index aabfa9a945a861a6475172cc262e3ba9ae42e3c0..77744ac52b58ec889945c8e47ea5717dfce5b6f5 100644 (file)
@@ -18,6 +18,7 @@ import os
 import json
 import logging
 import subprocess
+import pkg_resources
 
 from mediagoblin import mg_globals as mgg
 from mediagoblin.processing import create_pub_filepath, \
@@ -29,6 +30,17 @@ from mediagoblin.media_types.stl import model_loader
 _log = logging.getLogger(__name__)
 SUPPORTED_FILETYPES = ['stl', 'obj']
 
+BLEND_FILE = pkg_resources.resource_filename(
+    'mediagoblin.media_types.stl',
+    os.path.join(
+        'assets',
+        'blender_render.blend'))
+BLEND_SCRIPT = pkg_resources.resource_filename(
+    'mediagoblin.media_types.stl',
+    os.path.join(
+        'assets',
+        'blender_render.py'))
+
 
 def sniff_handler(media_file, **kw):
     if kw.get('media') is not None:
@@ -55,19 +67,23 @@ def blender_render(config):
     arg_string = "blender -b blender_render.blend -F "
     arg_string +="JPEG -P blender_render.py"
     env = {"RENDER_SETUP" : json.dumps(config), "DISPLAY":":0"}
-    subprocess.call(arg_string.split(" "), env=env)
+    subprocess.call(
+        ["blender",
+         "-b", BLEND_FILE,
+         "-F", "JPEG",
+         "-P", BLEND_SCRIPT],
+        env=env)
 
 
-def process_stl(entry):
-    """
-    Code to process an stl or obj model.
+def process_stl(proc_state):
+    """Code to process an stl or obj model. Will be run by celery.
+
+    A Workbench() represents a local tempory dir. It is automatically
+    cleaned up when this function exits.
     """
+    entry = proc_state.entry
+    workbench = proc_state.workbench
 
-    workbench = mgg.workbench_manager.create_workbench()
-    # Conversions subdirectory to avoid collisions
-    conversions_subdir = os.path.join(
-        workbench.dir, 'conversions')
-    os.mkdir(conversions_subdir)
     queued_filepath = entry.queued_media_file
     queued_filename = workbench.localized_file(
         mgg.queue_store, queued_filepath, 'source')
@@ -90,23 +106,33 @@ def process_stl(entry):
     greatest = greatest[-1]
 
     def snap(name, camera, width=640, height=640, project="ORTHO"):
-        path = create_pub_filepath(entry, name_builder.fill(name))
-        render_file = mgg.public_store.get_file(path, "wb")
+        filename = name_builder.fill(name)
+        workbench_path = workbench.joinpath(filename)
         shot = {
-            "model_path" : queued_filename,
-            "model_ext" : ext,
-            "camera_coord" : camera,
-            "camera_focus" : model.average,
-            "camera_clip" : greatest*10,
-            "greatest" : greatest,
-            "projection" : project,
-            "width" : width,
-            "height" : height,
-            "out_file" : render_file.name,
+            "model_path": queued_filename,
+            "model_ext": ext,
+            "camera_coord": camera,
+            "camera_focus": model.average,
+            "camera_clip": greatest*10,
+            "greatest": greatest,
+            "projection": project,
+            "width": width,
+            "height": height,
+            "out_file": workbench_path,
             }
-        render_file.close()
         blender_render(shot)
-        return path
+
+        # make sure the image rendered to the workbench path
+        assert os.path.exists(workbench_path)
+
+        # copy it up!
+        with open(workbench_path, 'rb') as rendered_file:
+            public_path = create_pub_filepath(entry, filename)
+
+            with mgg.public_store.get_file(public_path, "wb") as public_file:
+                public_file.write(rendered_file.read())
+
+        return public_path
 
     thumb_path = snap(
         "{basename}.thumb.jpg",
@@ -131,10 +157,7 @@ def process_stl(entry):
         "{basename}.side.jpg",
         [greatest*-2, model.average[1], model.average[2]])
 
-
-
-
-    # Save the public file stuffs
+    ## Save the public file stuffs
     model_filepath = create_pub_filepath(
         entry, name_builder.fill('{basename}{ext}'))
 
@@ -142,11 +165,10 @@ def process_stl(entry):
         with open(queued_filename, 'rb') as queued_file:
             model_file.write(queued_file.read())
 
-
     # Remove queued media file from storage and database
     mgg.queue_store.delete_file(queued_filepath)
     entry.queued_media_file = []
-        
+
     # Insert media file information into database
     media_files_dict = entry.setdefault('media_files', {})
     media_files_dict[u'original'] = model_filepath
@@ -164,8 +186,6 @@ def process_stl(entry):
         "width" : model.width,
         "height" : model.height,
         "depth" : model.depth,
+        "file_type" : ext,
         }
     entry.media_data_init(**dimensions)
-
-    # clean up workbench
-    workbench.destroy_self()