Better way to add optional extension to specified template file

nbconvert/exporters/templateexporter.py

@@ -14,7 +14,9 @@
 from traitlets import HasTraits, Unicode, List, Dict, default, observe
 from traitlets.utils.importstring import import_item
 from ipython_genutils import py3compat
-from jinja2 import TemplateNotFound, Environment, ChoiceLoader, FileSystemLoader
+from jinja2 import (
+    TemplateNotFound, Environment, ChoiceLoader, FileSystemLoader, BaseLoader
+)
 
 from nbconvert import filters
 from .exporter import Exporter
@@ -55,6 +57,31 @@
 }
 
 
+class ExtensionTolerantLoader(BaseLoader):
+    """A template loader which optionally adds a given extension when searching.
+
+    Constructor takes two arguments: *loader* is another Jinja loader instance
+    to wrap. *extension* is the extension, which will be added to the template
+    name if finding the template without it fails. This should include the dot,
+    e.g. '.tpl'.
+    """
+    def __init__(self, loader, extension):
+        self.loader = loader
+        self.extension = extension
+
+
+    def get_source(self, environment, template):
+        try:
+            return self.loader.get_source(environment, template)
+        except TemplateNotFound:
+            if template.endswith(self.extension):
+                raise TemplateNotFound(template)
+            return self.loader.get_source(environment, template+self.extension)
+
+    def list_templates(self):
+        return self.loader.list_templates()
+
+
 class TemplateExporter(Exporter):
     """
     Exports notebooks into other file formats.  Uses Jinja 2 templating engine
@@ -135,7 +162,7 @@ def _template_file_default(self):
     ).tag(affects_environment=True)
 
     #Extension that the template files use.
-    template_extension = Unicode(".tpl").tag(config=True, affects_template=True)
+    template_extension = Unicode(".tpl").tag(config=True, affects_environment=True)
 
     extra_loaders = List(
         help="Jinja loaders to find templates. Will be tried in order "
@@ -190,19 +217,9 @@ def _load_template(self):
         # template by name with extension added, then try loading the template
         # as if the name is explicitly specified.
         template_file = self.template_file
-        if not template_file.endswith(self.template_extension):
-            template_file = template_file + self.template_extension
         self.log.debug("Attempting to load template %s", template_file)
         self.log.debug("    template_path: %s", os.pathsep.join(self.template_path))
-        try:
-            template = self.environment.get_template(template_file)
-        except (TemplateNotFound, IOError):
-            pass
-        else:
-            self.log.debug("Loaded template %s", template.filename)
-            return template
-
-        raise TemplateNotFound(template_file)
+        return self.environment.get_template(template_file)
 
     def from_notebook_node(self, nb, resources=None, **kw):
         """
@@ -303,7 +320,9 @@ def _create_environment(self):
             [os.path.join(here, self.default_template_path),
              os.path.join(here, self.template_skeleton_path)]
 
-        loaders = self.extra_loaders + [FileSystemLoader(paths)]
+        loaders = self.extra_loaders + [
+            ExtensionTolerantLoader(FileSystemLoader(paths), self.template_extension)
+        ]
         environment = Environment(
             loader=ChoiceLoader(loaders),
             extensions=JINJA_EXTENSIONS

nbconvert/exporters/tests/test_latex.py

@@ -14,6 +14,7 @@
 from ipython_genutils.testing.decorators import onlyif_cmds_exist
 from testpath.tempdir import TemporaryDirectory
 
+from jinja2 import DictLoader
 
 class TestLatexExporter(ExportersTestsBase):
     """Contains test functions for latex.py"""
@@ -115,3 +116,16 @@ def test_prompt_number_color(self):
 
         assert re.findall(in_regex, output) == ins
         assert re.findall(out_regex, output) == outs
+
+    def test_in_memory_template_tplx(self):
+        # Loads in an in memory latex template (.tplx) using jinja2.DictLoader
+        # creates a class that uses this template with the template_file argument
+        # converts an empty notebook using this mechanism
+        my_loader_tplx = DictLoader({'my_template': "{%- extends 'article.tplx' -%}"})
+
+        class MyExporter(LatexExporter):
+            template_file = 'my_template'
+
+        exporter = MyExporter(extra_loaders=[my_loader_tplx])
+        nb = v4.new_notebook()
+        out, resources = exporter.from_notebook_node(nb)

nbconvert/exporters/tests/test_templateexporter.py

@@ -8,12 +8,15 @@
 import os
 
 from traitlets.config import Config
+from jinja2 import DictLoader, TemplateNotFound
+from nbformat import v4
 
 from .base import ExportersTestsBase
 from .cheese import CheesePreprocessor
 from ..templateexporter import TemplateExporter
 from testpath import tempdir
 
+import pytest
 
 class TestExporter(ExportersTestsBase):
     """Contains test functions for exporter.py"""
@@ -114,6 +117,35 @@ def test_relative_template_file(self):
             exporter = self._make_exporter(config=config)
             assert os.path.abspath(exporter.template.filename) == template
             assert os.path.dirname(template) in [ os.path.abspath(d) for d in exporter.template_path ]
+
+    def test_in_memory_template(self):
+        # Loads in an in memory template using jinja2.DictLoader
+        # creates a class that uses this template with the template_file argument
+        # converts an empty notebook using this mechanism
+        my_loader = DictLoader({'my_template': "{%- extends 'rst.tpl' -%}"})
+
+        class MyExporter(TemplateExporter):
+            template_file = 'my_template'
+
+        exporter = MyExporter(extra_loaders=[my_loader])
+        nb = v4.new_notebook()
+        out, resources = exporter.from_notebook_node(nb)
+
+
+    @pytest.mark.xfail(strict=True, raises=TemplateNotFound)
+    def test_in_memory_template_failure_to_find(self):
+
+        # Create exporter with invalid template file, try to convert empty notebook
+        # failure is expected due to nonexistant template file.
+
+        template = 'does_not_exist.tpl'
+        exporter = TemplateExporter(template_file=template)
+        assert template not in exporter.environment.loader.list_templates()
+        nb = v4.new_notebook()
+        out, resources = exporter.from_notebook_node(nb)
+
+
+
 
     def _make_exporter(self, config=None):
         # Create the exporter instance, make sure to set a template name since