mediagoblin/tools/workbench.py

   1 # GNU MediaGoblin -- federated, autonomous media hosting
   2 # Copyright (C) 2011, 2012 MediaGoblin contributors.  See AUTHORS.
   3 #
   4 # This program is free software: you can redistribute it and/or modify
   5 # it under the terms of the GNU Affero General Public License as published by
   6 # the Free Software Foundation, either version 3 of the License, or
   7 # (at your option) any later version.
   8 #
   9 # This program is distributed in the hope that it will be useful,
  10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 # GNU Affero General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU Affero General Public License
  15 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16
  17 import os
  18 import shutil
  19 import tempfile
  20
  21 import six
  22
  23 from mediagoblin._compat import py2_unicode
  24
  25 # Actual workbench stuff
  26 # ----------------------
  27
  28
  29 @py2_unicode
  30 class Workbench(object):
  31     """
  32     Represent the directory for the workbench
  33
  34     WARNING: DO NOT create Workbench objects on your own,
  35     let the WorkbenchManager do that for you!
  36     """
  37     def __init__(self, dir):
  38         """
  39         WARNING: DO NOT create Workbench objects on your own,
  40         let the WorkbenchManager do that for you!
  41         """
  42         self.dir = dir
  43
  44     def __str__(self):
  45         return six.text_type(self.dir)
  46
  47     def __repr__(self):
  48         try:
  49             return str(self)
  50         except AttributeError:
  51             return 'None'
  52
  53     def joinpath(self, *args):
  54         return os.path.join(self.dir, *args)
  55
  56     def localized_file(self, storage, filepath,
  57                        filename_if_copying=None,
  58                        keep_extension_if_copying=True):
  59         """
  60         Possibly localize the file from this storage system (for read-only
  61         purposes, modifications should be written to a new file.).
  62
  63         If the file is already local, just return the absolute filename of that
  64         local file.  Otherwise, copy the file locally to the workbench, and
  65         return the absolute path of the new file.
  66
  67         If it is copying locally, we might want to require a filename like
  68         "source.jpg" to ensure that we won't conflict with other filenames in
  69         our workbench... if that's the case, make sure filename_if_copying is
  70         set to something like 'source.jpg'.  Relatedly, if you set
  71         keep_extension_if_copying, you don't have to set an extension on
  72         filename_if_copying yourself, it'll be set for you (assuming such an
  73         extension can be extacted from the filename in the filepath).
  74
  75         Returns:
  76           localized_filename
  77
  78         Examples:
  79           >>> wb_manager.localized_file(
  80           ...     '/our/workbench/subdir', local_storage,
  81           ...     ['path', 'to', 'foobar.jpg'])
  82           u'/local/storage/path/to/foobar.jpg'
  83
  84           >>> wb_manager.localized_file(
  85           ...     '/our/workbench/subdir', remote_storage,
  86           ...     ['path', 'to', 'foobar.jpg'])
  87           '/our/workbench/subdir/foobar.jpg'
  88
  89           >>> wb_manager.localized_file(
  90           ...     '/our/workbench/subdir', remote_storage,
  91           ...     ['path', 'to', 'foobar.jpg'], 'source.jpeg', False)
  92           '/our/workbench/subdir/foobar.jpeg'
  93
  94           >>> wb_manager.localized_file(
  95           ...     '/our/workbench/subdir', remote_storage,
  96           ...     ['path', 'to', 'foobar.jpg'], 'source', True)
  97           '/our/workbench/subdir/foobar.jpg'
  98         """
  99         if storage.local_storage:
 100             return storage.get_local_path(filepath)
 101         else:
 102             if filename_if_copying is None:
 103                 dest_filename = filepath[-1]
 104             else:
 105                 orig_filename, orig_ext = os.path.splitext(filepath[-1])
 106                 if keep_extension_if_copying and orig_ext:
 107                     dest_filename = filename_if_copying + orig_ext
 108                 else:
 109                     dest_filename = filename_if_copying
 110
 111             full_dest_filename = os.path.join(
 112                 self.dir, dest_filename)
 113
 114             # copy it over
 115             storage.copy_locally(
 116                 filepath, full_dest_filename)
 117
 118             return full_dest_filename
 119
 120     def destroy(self):
 121         """
 122         Destroy this workbench!  Deletes the directory and all its contents!
 123
 124         WARNING: Does no checks for a sane value in self.dir!
 125         """
 126         # just in case
 127         workbench = os.path.abspath(self.dir)
 128         shutil.rmtree(workbench)
 129         del self.dir
 130
 131     def __enter__(self):
 132         """Make Workbench a context manager so we can use `with Workbench() as bench:`"""
 133         return self
 134
 135     def  __exit__(self, *args):
 136         """Clean up context manager, aka ourselves, deleting the workbench"""
 137         self.destroy()
 138
 139
 140 class WorkbenchManager(object):
 141     """
 142     A system for generating and destroying workbenches.
 143
 144     Workbenches are actually just subdirectories of a (local) temporary
 145     storage space for during the processing stage. The preferred way to
 146     create them is to use:
 147
 148         with workbenchmger.create() as workbench:
 149             do stuff...
 150
 151     This will automatically clean up all temporary directories even in
 152     case of an exceptions. Also check the
 153     @mediagoblin.decorators.get_workbench decorator for a convenient
 154     wrapper.
 155     """
 156
 157     def __init__(self, base_workbench_dir):
 158         self.base_workbench_dir = os.path.abspath(base_workbench_dir)
 159         if not os.path.exists(self.base_workbench_dir):
 160             os.makedirs(self.base_workbench_dir)
 161
 162     def create(self):
 163         """
 164         Create and return the path to a new workbench (directory).
 165         """
 166         return Workbench(tempfile.mkdtemp(dir=self.base_workbench_dir))