use name_builder with store_public, not create_pub_filepath
[mediagoblin.git] / mediagoblin / media_types / audio / processing.py
index f0b8d0f9c2b957bbd88d3ff44d1587fd355f8a2a..6c565eb4c69974c011525ecaa2f4f3ca5a614868 100644 (file)
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import argparse
 import logging
-import tempfile
+from tempfile import NamedTemporaryFile
 import os
 
 from mediagoblin import mg_globals as mgg
-from mediagoblin.processing import create_pub_filepath, BadMediaFail
+from mediagoblin.processing import (
+    create_pub_filepath, BadMediaFail, FilenameBuilder,
+    ProgressCallback, MediaProcessor, ProcessingManager,
+    request_from_args, get_orig_filename,
+    store_public, copy_original)
 
-from mediagoblin.media_types.audio.transcoders import AudioTranscoder, \
-    AudioThumbnailer
+from mediagoblin.media_types.audio.transcoders import (
+    AudioTranscoder,AudioThumbnailer)
 
 _log = logging.getLogger(__name__)
 
+MEDIA_TYPE = 'mediagoblin.media_types.audio'
+
+
 def sniff_handler(media_file, **kw):
+    _log.info('Sniffing {0}'.format(MEDIA_TYPE))
     try:
         transcoder = AudioTranscoder()
         data = transcoder.discover(media_file.name)
     except BadMediaFail:
         _log.debug('Audio discovery raised BadMediaFail')
-        return False
+        return None
 
     if data.is_audio == True and data.is_video == False:
-        return True
+        return MEDIA_TYPE
 
-    return False
+    return None
 
-def process_audio(entry):
-    audio_config = mgg.global_config['media_type:mediagoblin.media_types.audio']
 
-    workbench = mgg.workbench_manager.create_workbench()
+def process_audio(proc_state):
+    """Code to process uploaded audio. 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
+    audio_config = mgg.global_config['media_type:mediagoblin.media_types.audio']
 
     queued_filepath = entry.queued_media_file
     queued_filename = workbench.localized_file(
         mgg.queue_store, queued_filepath,
         'source')
+    name_builder = FilenameBuilder(queued_filename)
 
-    ogg_filepath = create_pub_filepath(
+    webm_audio_filepath = create_pub_filepath(
         entry,
         '{original}.webm'.format(
             original=os.path.splitext(
                 queued_filepath[-1])[0]))
 
+    if audio_config['keep_original']:
+        with open(queued_filename, 'rb') as queued_file:
+            original_filepath = create_pub_filepath(
+                entry, name_builder.fill('{basename}{ext}'))
+
+            with mgg.public_store.get_file(original_filepath, 'wb') as \
+                    original_file:
+                _log.debug('Saving original...')
+                original_file.write(queued_file.read())
+
+            entry.media_files['original'] = original_filepath
+
     transcoder = AudioTranscoder()
 
-    with tempfile.NamedTemporaryFile() as ogg_tmp:
+    with NamedTemporaryFile(dir=workbench.dir) as webm_audio_tmp:
+        progress_callback = ProgressCallback(entry)
 
         transcoder.transcode(
             queued_filename,
-            ogg_tmp.name,
-            quality=audio_config['quality'])
+            webm_audio_tmp.name,
+            quality=audio_config['quality'],
+            progress_callback=progress_callback)
 
-        data = transcoder.discover(ogg_tmp.name)
+        transcoder.discover(webm_audio_tmp.name)
 
         _log.debug('Saving medium...')
-        mgg.public_store.get_file(ogg_filepath, 'wb').write(
-            ogg_tmp.read())
+        mgg.public_store.get_file(webm_audio_filepath, 'wb').write(
+            webm_audio_tmp.read())
 
-        entry.media_files['ogg'] = ogg_filepath
+        entry.media_files['webm_audio'] = webm_audio_filepath
 
         # entry.media_data_init(length=int(data.audiolength))
 
@@ -81,16 +111,17 @@ def process_audio(entry):
                 original=os.path.splitext(
                     queued_filepath[-1])[0]))
 
-        with tempfile.NamedTemporaryFile(suffix='.wav') as wav_tmp:
-            _log.info('Creating WAV source for spectrogram')
+        with NamedTemporaryFile(dir=workbench.dir, suffix='.ogg') as wav_tmp:
+            _log.info('Creating OGG source for spectrogram')
             transcoder.transcode(
                 queued_filename,
                 wav_tmp.name,
-                mux_string='wavenc')
+                mux_string='vorbisenc quality={0} ! oggmux'.format(
+                    audio_config['quality']))
 
             thumbnailer = AudioThumbnailer()
 
-            with tempfile.NamedTemporaryFile(suffix='.jpg') as spectrogram_tmp:
+            with NamedTemporaryFile(dir=workbench.dir, suffix='.jpg') as spectrogram_tmp:
                 thumbnailer.spectrogram(
                     wav_tmp.name,
                     spectrogram_tmp.name,
@@ -103,7 +134,7 @@ def process_audio(entry):
 
                 entry.media_files['spectrogram'] = spectrogram_filepath
 
-                with tempfile.NamedTemporaryFile(suffix='.jpg') as thumb_tmp:
+                with NamedTemporaryFile(dir=workbench.dir, suffix='.jpg') as thumb_tmp:
                     thumbnailer.thumbnail_spectrogram(
                         spectrogram_tmp.name,
                         thumb_tmp.name,
@@ -123,9 +154,201 @@ def process_audio(entry):
     else:
         entry.media_files['thumb'] = ['fake', 'thumb', 'path.jpg']
 
-    mgg.queue_store.delete_file(queued_filepath)
+    # Remove queued media file from storage and database.
+    # queued_filepath is in the task_id directory which should
+    # be removed too, but fail if the directory is not empty to be on
+    # the super-safe side.
+    mgg.queue_store.delete_file(queued_filepath)      # rm file
+    mgg.queue_store.delete_dir(queued_filepath[:-1])  # rm dir
+    entry.queued_media_file = []
+
+
+class CommonAudioProcessor(MediaProcessor):
+    """
+    Provides a base for various audio processing steps
+    """
+
+    def common_setup(self):
+        """
+        """
+        self.audio_config = mgg \
+            .global_config['media_type:mediagoblin.media_types.audio']
+
+        # Pull down and set up the original file
+        self.orig_filename = get_orig_filename(
+            self.entry, self.workbench)
+        self.name_builder = FilenameBuilder(self.orig_filename)
+
+        self.spectrogram_tmp = os.path.join(self.workbench.dir,
+                                            self.name_builder.fill(
+                                                '{basename}-spectrogram.jpg'))
+
+        self.transcoder = AudioTranscoder()
+        self.thumbnailer = AudioThumbnailer()
+
+    def copy_original(self):
+        if self.audio_config['keep_original']:
+            copy_original(
+                self.entry, self.orig_filename,
+                self.name_builder.fill('{basename}{ext}'))
+
+    def transcode(self, quality=None):
+        if not quality:
+            quality = self.audio_config['quality']
+
+        progress_callback = ProgressCallback(self.entry)
+        webm_audio_tmp = os.path.join(self.workbench.dir,
+                                      self.name_builder.fill(
+                                          '{basename}{ext}'))
+
+        #webm_audio_filepath = create_pub_filepath(
+        #    self.entry,
+        #    '{original}.webm'.format(
+        #        original=os.path.splitext(
+        #            self.orig_filename[-1])[0]))
+
+        self.transcoder.transcode(
+            self.orig_filename,
+            webm_audio_tmp,
+            quality=quality,
+            progress_callback=progress_callback)
+
+        self.transcoder.discover(webm_audio_tmp)
+
+        _log.debug('Saving medium...')
+        store_public(self.entry, 'medium', webm_audio_tmp,
+                     self.name_builder.fill('{basename}.medium{ext}'))
+
+    def create_spectrogram(self, quality=None, max_width=None, fft_size=None):
+        if not quality:
+            quality = self.audio_config['quality']
+        if not max_width:
+            max_width = mgg.global_config['media:medium']['max_width']
+        if not fft_size:
+            fft_size = self.audio_config['spectrogram_fft_size']
+
+        #spectrogram_filepath = create_pub_filepath(
+        #    self.entry,
+        #    '{original}-spectrogram.jpg'.format(
+        #        original=os.path.splitext(
+        #            self.orig_filename[-1])[0]))
+
+        wav_tmp = os.path.join(self.workbench.dir, self.name_builder.fill(
+            '{basename}.ogg'))
+
+        _log.info('Creating OGG source for spectrogram')
+        self.transcoder.transcode(
+            self.orig_filename,
+            wav_tmp,
+            mux_string='vorbisenc quality={0} ! oggmux'.format(quality))
+
+        self.thumbnailer.spectrogram(
+            wav_tmp,
+            self.spectrogram_tmp,
+            width=max_width,
+            fft_size=fft_size)
+
+        _log.debug('Saving spectrogram...')
+        store_public(self.entry, 'spectrogram', self.spectrogram_tmp,
+                     self.name_builder.fill('{basename}.spectrogram.jpg'))
+
+    def generate_thumb(self, size=None):
+        if not size:
+            max_width = mgg.global_config['medium:thumb']['max_width']
+            max_height = mgg.global_config['medium:thumb']['max_height']
+            size = (max_width, max_height)
+
+        thumb_tmp = os.path.join(self.workbench.dir, self.name_builder.fill(
+            '{basename}-thumbnail.jpg'))
+
+        self.thumbnailer.thumbnail_spectrogram(
+            self.spectrogram_tmp,
+            thumb_tmp,
+            size)
+
+        #thumb_filepath = create_pub_filepath(
+        #    self.entry,
+        #    '{original}-thumbnail.jpg'.format(
+        #        original=os.path.splitext(
+        #            self.orig_filename[-1])[0]))
+
+        store_public(self.entry, 'thumb', thumb_tmp,
+                     self.name_builder.fill('{basename}.thumbnail.jpg'))
+
+
+class InitialProcessor(CommonAudioProcessor):
+    """
+    Initial processing steps for new audio
+    """
+    name = "initial"
+    description = "Initial processing"
+
+    @classmethod
+    def media_is_eligible(cls, entry=None, state=None):
+        """
+        Determine if this media type is eligible for processing
+        """
+        if not state:
+            state = entry.state
+        return state in (
+            "unprocessed", "failed")
+
+    @classmethod
+    def generate_parser(cls):
+        parser = argparse.ArgumentParser(
+            description=cls.description,
+            prog=cls.name)
+
+        parser.add_argument(
+            '--quality',
+            help='vorbisenc quality')
+
+        parser.add_argument(
+            '--fft_size',
+            type=int,
+            help='spectrogram fft size')
+
+        parser.add_argument(
+            '--thumb_size',
+            metavar=('max_width', 'max_height'),
+            type=int)
+
+        parser.add_argument(
+            '--medium_width',
+            type=int,
+            help='The width of the spectogram')
+
+        parser.add_argument(
+            '--create_spectrogram',
+            action='store_true',
+            help='Create spectogram and thumbnail')
+
+        return parser
+
+    @classmethod
+    def args_to_request(cls, args):
+        return request_from_args(
+            args, ['create_spectrogram', 'quality', 'fft_size',
+                   'thumb_size', 'medium_width'])
+
+    def process(self, quality=None, fft_size=None, thumb_size=None,
+                create_spectrogram=None, medium_width=None):
+        self.common_setup()
+
+        if not create_spectrogram:
+            create_spectrogram = self.audio_config['create_spectrogram']
+
+        self.transcode(quality=quality)
+        self.copy_original()
+
+        if create_spectrogram:
+            self.create_spectrogram(quality=quality, max_width=medium_width,
+                                    fft_size=fft_size)
+            self.generate_thumb(size=thumb_size)
+        self.delete_queue_file()
 
-    entry.save()
 
-    # clean up workbench
-    workbench.destroy_self()
+class AudioProcessingManager(ProcessingManager):
+    def __init__(self):
+        super(self.__class__, self).__init__()
+        self.add_processor(InitialProcessor)