From 6fa6cb0e129893b0166448563597e1d156cd8746 Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Tue, 13 Sep 2011 23:27:41 +0200
Subject: [PATCH] Improved choice and prefix loaders.

Choice and prefix loaders now dispatch source and template lookup
separately in order to work in combination with module loaders as
advertised.
---
 CHANGES                    |  8 ++++++++
 jinja2/loaders.py          | 25 ++++++++++++++++++++++++-
 jinja2/testsuite/loader.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/CHANGES b/CHANGES
index ae07828..4427436 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,6 +1,14 @@
 Jinja2 Changelog
 ================
 
+Version 2.7
+-----------
+(codename to be selected, release date to be announced)
+
+- Choice and prefix loaders now dispatch source and template lookup
+  separately in order to work in combination with module loaders as
+  advertised.
+
 Version 2.6
 -----------
 (codename Convolution, released on July 24th 2011)
diff --git a/jinja2/loaders.py b/jinja2/loaders.py
index 419a9c8..c90bbe7 100644
--- a/jinja2/loaders.py
+++ b/jinja2/loaders.py
@@ -330,12 +330,16 @@ class PrefixLoader(BaseLoader):
         self.mapping = mapping
         self.delimiter = delimiter
 
-    def get_source(self, environment, template):
+    def get_loader(self, template):
         try:
             prefix, name = template.split(self.delimiter, 1)
             loader = self.mapping[prefix]
         except (ValueError, KeyError):
             raise TemplateNotFound(template)
+        return loader, name
+
+    def get_source(self, environment, template):
+        loader, name = self.get_loader(template)
         try:
             return loader.get_source(environment, name)
         except TemplateNotFound:
@@ -343,6 +347,16 @@ class PrefixLoader(BaseLoader):
             # (the one that includes the prefix)
             raise TemplateNotFound(template)
 
+    @internalcode
+    def load(self, environment, name, globals=None):
+        loader, local_name = self.get_loader(name)
+        try:
+            return loader.load(environment, local_name)
+        except TemplateNotFound:
+            # re-raise the exception with the correct fileame here.
+            # (the one that includes the prefix)
+            raise TemplateNotFound(name)
+
     def list_templates(self):
         result = []
         for prefix, loader in self.mapping.iteritems():
@@ -376,6 +390,15 @@ class ChoiceLoader(BaseLoader):
                 pass
         raise TemplateNotFound(template)
 
+    @internalcode
+    def load(self, environment, name, globals=None):
+        for loader in self.loaders:
+            try:
+                return loader.load(environment, name, globals)
+            except TemplateNotFound:
+                pass
+        raise TemplateNotFound(name)
+
     def list_templates(self):
         found = set()
         for loader in self.loaders:
diff --git a/jinja2/testsuite/loader.py b/jinja2/testsuite/loader.py
index fb1e53d..f62ec92 100644
--- a/jinja2/testsuite/loader.py
+++ b/jinja2/testsuite/loader.py
@@ -182,6 +182,34 @@ class ModuleLoaderTestCase(JinjaTestCase):
             tmpl_3c4ddf650c1a73df961a6d3d2ce2752f1b8fd490
         assert mod.__file__.endswith('.pyc')
 
+    def test_choice_loader(self):
+        log = self.compile_down(py_compile=True)
+        assert 'Byte-compiled "a/test.html"' in log
+
+        self.mod_env.loader = loaders.ChoiceLoader([
+            self.mod_env.loader,
+            loaders.DictLoader({'DICT_SOURCE': 'DICT_TEMPLATE'})
+        ])
+
+        tmpl1 = self.mod_env.get_template('a/test.html')
+        self.assert_equal(tmpl1.render(), 'BAR')
+        tmpl2 = self.mod_env.get_template('DICT_SOURCE')
+        self.assert_equal(tmpl2.render(), 'DICT_TEMPLATE')
+
+    def test_prefix_loader(self):
+        log = self.compile_down(py_compile=True)
+        assert 'Byte-compiled "a/test.html"' in log
+
+        self.mod_env.loader = loaders.PrefixLoader({
+            'MOD':      self.mod_env.loader,
+            'DICT':     loaders.DictLoader({'test.html': 'DICT_TEMPLATE'})
+        })
+
+        tmpl1 = self.mod_env.get_template('MOD/a/test.html')
+        self.assert_equal(tmpl1.render(), 'BAR')
+        tmpl2 = self.mod_env.get_template('DICT/test.html')
+        self.assert_equal(tmpl2.render(), 'DICT_TEMPLATE')
+
 
 def suite():
     suite = unittest.TestSuite()
-- 
2.26.2