- [feature] Added module_writer argument to Template,

  TemplateLookup, allows a callable to be passed which
  takes over the writing of the template's module source
  file, so that special environment-specific steps
  can be taken.  [ticket:181]

4 files changed, 81 insertions, 12 deletions

diff --git a/CHANGES b/CHANGES
index a8455fd..388c9b7 100644
--- a/CHANGES
+++ b/CHANGES
@@ -37,6 +37,12 @@
   script, allows passing of kw to the template from 
   the command line. [ticket:178]
 
+- [feature] Added module_writer argument to Template, 
+  TemplateLookup, allows a callable to be passed which
+  takes over the writing of the template's module source
+  file, so that special environment-specific steps
+  can be taken.  [ticket:181]
+
 - [bug] The exception message in the html_error_template
   is now escaped with the HTML filter. [ticket:142]
 
diff --git a/mako/lookup.py b/mako/lookup.py
index ef75ade..f6f811b 100644
--- a/mako/lookup.py
+++ b/mako/lookup.py
@@ -160,6 +160,7 @@ class TemplateLookup(TemplateCollection):
                         cache_url=None, 
 
                         modulename_callable=None, 
+                        module_writer=None,
                         default_filters=None, 
                         buffer_filters=(), 
                         strict_undefined=False,
@@ -195,6 +196,7 @@ class TemplateLookup(TemplateCollection):
             'encoding_errors':encoding_errors, 
             'input_encoding':input_encoding, 
             'module_directory':module_directory, 
+            'module_writer':module_writer,
             'cache_args':cache_args,
             'cache_enabled':cache_enabled, 
             'default_filters':default_filters, 
diff --git a/mako/template.py b/mako/template.py
index 5d30c13..dcfc9f1 100644
--- a/mako/template.py
+++ b/mako/template.py
@@ -110,6 +110,39 @@ class Template(object):
     :param module_filename: Overrides the filename of the generated 
      Python module file. For advanced usage only.
  
+    :param module_writer: A callable which overrides how the Python
+     module is written entirely.  The callable is passed the 
+     encoded source content of the module and the destination
+     path to be written to.   The default behavior of module writing
+     uses a tempfile in conjunction with a file move in order
+     to make the operation atomic.   So a user-defined module
+     writing function that mimics the default behavior would be::
+
+        import tempfile
+        import os
+        import shutil
+
+        def module_writer(source, outputpath):
+            (dest, name) = \\
+                tempfile.mkstemp(
+                    dir=os.path.dirname(outputpath)
+                )
+ 
+            os.write(dest, source)
+            os.close(dest)
+            shutil.move(name, outputpath)
+
+        from mako.template import Template
+        mytemplate = Template(
+                        file="index.html", 
+                        module_directory="/path/to/modules", 
+                        module_writer=module_writer
+                    )
+
+     The function is provided for unusual configurations where
+     certain platform-specific permissions or other special
+     steps are needed.
+
     :param output_encoding: The encoding to use when :meth:`.render` 
      is called.  
      See :ref:`usage_unicode` as well as :ref:`unicode_toplevel`.
@@ -154,6 +187,7 @@ class Template(object):
                     module_filename=None, 
                     input_encoding=None, 
                     disable_unicode=False,
+                    module_writer=None,
                     bytestring_passthrough=False, 
                     default_filters=None, 
                     buffer_filters=(), 
@@ -188,6 +222,7 @@ class Template(object):
         self.disable_unicode = disable_unicode
         self.bytestring_passthrough = bytestring_passthrough or disable_unicode
         self.strict_undefined = strict_undefined
+        self.module_writer = module_writer
 
         if util.py3k and disable_unicode:
             raise exceptions.UnsupportedError(
@@ -276,7 +311,8 @@ class Template(object):
                             self, 
                             open(filename, 'rb').read(), 
                             filename, 
-                            path)
+                            path,
+                            self.module_writer)
             module = imp.load_source(self.module_id, path, open(path, 'rb'))
             del sys.modules[self.module_id]
             if module._magic_number != codegen.MAGIC_NUMBER:
@@ -284,7 +320,8 @@ class Template(object):
                             self, 
                             open(filename, 'rb').read(), 
                             filename, 
-                            path)
+                            path,
+                            self.module_writer)
                 module = imp.load_source(self.module_id, path, open(path, 'rb'))
                 del sys.modules[self.module_id]
             ModuleInfo(module, path, self, filename, None, None)
@@ -543,7 +580,7 @@ def _compile_text(template, text, filename):
     exec code in module.__dict__, module.__dict__
     return (source, module)
 
-def _compile_module_file(template, text, filename, outputpath):
+def _compile_module_file(template, text, filename, outputpath, module_writer):
     identifier = template.module_id
     lexer = Lexer(text, 
                     filename, 
@@ -563,17 +600,20 @@ def _compile_module_file(template, text, filename, outputpath):
                                 disable_unicode=template.disable_unicode,
                                 strict_undefined=template.strict_undefined)
  
-    # make tempfiles in the same location as the ultimate 
-    # location.   this ensures they're on the same filesystem,
-    # avoiding synchronization issues.
-    (dest, name) = tempfile.mkstemp(dir=os.path.dirname(outputpath))
- 
     if isinstance(source, unicode):
         source = source.encode(lexer.encoding or 'ascii')
- 
-    os.write(dest, source)
-    os.close(dest)
-    shutil.move(name, outputpath)
+
+    if module_writer:
+        module_writer(source, outputpath)
+    else:
+        # make tempfiles in the same location as the ultimate 
+        # location.   this ensures they're on the same filesystem,
+        # avoiding synchronization issues.
+        (dest, name) = tempfile.mkstemp(dir=os.path.dirname(outputpath))
+ 
+        os.write(dest, source)
+        os.close(dest)
+        shutil.move(name, outputpath)
 
 def _get_module_info_from_callable(callable_):
     return _get_module_info(callable_.func_globals['__name__'])
diff --git a/test/test_template.py b/test/test_template.py
index ddab16e..769af02 100644
--- a/test/test_template.py
+++ b/test/test_template.py
@@ -841,6 +841,10 @@ for utf8 in (True, False):
             del _do_test
 
 class ModuleDirTest(TemplateTest):
+    def tearDown(self):
+        import shutil
+        shutil.rmtree(module_base, True)
+
     def test_basic(self):
         t = self._file_template("modtest.html")
         t2 = self._file_template('subdir/modtest.html')
@@ -874,6 +878,22 @@ class ModuleDirTest(TemplateTest):
             os.path.join(module_base, 'subdir', 'foo', 'modtest.html.py')
         )
 
+    def test_custom_writer(self):
+        canary = []
+        def write_module(source, outputpath):
+            with open(outputpath, 'wb') as f:
+                canary.append(outputpath)
+                f.write(source)
+        lookup = TemplateLookup(template_base, module_writer=write_module, 
+                                            module_directory=module_base)
+        t = lookup.get_template('/modtest.html')
+        t2 = lookup.get_template('/subdir/modtest.html')
+        eq_(
+            canary,
+            [os.path.join(module_base, "modtest.html.py"),
+            os.path.join(module_base, "subdir/modtest.html.py")]
+        )
+
 class FilenameToURITest(TemplateTest):
     def test_windows_paths(self):
         """test that windows filenames are handled appropriately by Template."""
@@ -954,6 +974,7 @@ class FilenameToURITest(TemplateTest):
         t = Template("test", uri="foo/bar/../../foo.html")
         eq_(t.uri, "foo/bar/../../foo.html")
 
+
 class ModuleTemplateTest(TemplateTest):
     def test_module_roundtrip(self):
         lookup = TemplateLookup()