Correctly set template.self in all cases

3 files changed, 38 insertions, 16 deletions

diff --git a/lib/output/website/index.js b/lib/output/website/index.js
index 9ef87b6..0a8618c 100644
--- a/lib/output/website/index.js
+++ b/lib/output/website/index.js
@@ -154,12 +154,6 @@ WebsiteOutput.prototype.renderAsString = function(tpl, context) {
     // Calcul template name
     var filename = this.templateName(tpl);
 
-    // Extend context
-    // Setup complete context
-    context.template = _.extend(context.template || {}, {
-        self: filename
-    });
-
     context = _.extend(context, {
         plugins: {
             resources: this.resources
diff --git a/lib/output/website/templateEnv.js b/lib/output/website/templateEnv.js
index ea2b521..b3009db 100644
--- a/lib/output/website/templateEnv.js
+++ b/lib/output/website/templateEnv.js
@@ -2,6 +2,8 @@ var _ = require('lodash');
 var nunjucks = require('nunjucks');
 var path = require('path');
 var fs = require('fs');
+var DoExtension = require('nunjucks-do')(nunjucks);
+
 
 var location = require('../../utils/location');
 var defaultFilters = require('../../template/filters');
@@ -24,6 +26,8 @@ function setupTemplateEnv(output, context) {
     );
     var env = new nunjucks.Environment(loader);
 
+    env.addExtension('DoExtension', new DoExtension());
+
     // Add context as global
     _.each(context, function(value, key) {
         env.addGlobal(key, value);
diff --git a/lib/output/website/themeLoader.js b/lib/output/website/themeLoader.js
index 013b81c..bcfea5a 100644
--- a/lib/output/website/themeLoader.js
+++ b/lib/output/website/themeLoader.js
@@ -29,13 +29,16 @@ var ThemeLoader = nunjucks.Loader.extend({
         if (!fullpath) return null;
 
         fullpath = this.resolve(null, fullpath);
+        var templateName = this.getTemplateName(fullpath);
 
         if(!fullpath) {
             return null;
         }
 
         return {
-            src: fs.readFileSync(fullpath, 'utf-8'),
+            src: '{% do %}template = template || {}; template.stack = template.stack || []; template.stack.push(template.self); template.self = ' + JSON.stringify(templateName) + '{% enddo %}\n' +
+                fs.readFileSync(fullpath, 'utf-8') +
+                '\n{% do %}template.self = template.stack.pop();{% enddo %}',
             path: fullpath,
             noCache: true
         };
@@ -50,6 +53,34 @@ var ThemeLoader = nunjucks.Loader.extend({
     },
 
     /*
+        Get original search path containing a template
+
+        @param {String} filepath
+        @return {String} searchPath
+    */
+    getSearchPath: function(filepath) {
+        return _.chain(this.searchPaths)
+            .sortBy(function(s) {
+                return -s.length;
+            })
+            .find(function(basePath) {
+                return (filepath && filepath.indexOf(basePath) === 0);
+            })
+            .value();
+    },
+
+    /*
+        Get template name from a filepath
+
+        @param {String} filepath
+        @return {String} name
+    */
+    getTemplateName: function(filepath) {
+        var originalSearchPath = this.getSearchPath(filepath);
+        return originalSearchPath? path.relative(originalSearchPath, filepath) : null;
+    },
+
+    /*
         Resolve a template from a current template
 
         @param {String|null} from
@@ -65,15 +96,8 @@ var ThemeLoader = nunjucks.Loader.extend({
         }
 
         // Determine in which search folder we currently are
-        var originalSearchPath = _.chain(this.searchPaths)
-            .sortBy(function(s) {
-                return -s.length;
-            })
-            .find(function(basePath) {
-                return (from && from.indexOf(basePath) === 0);
-            })
-            .value();
-        var originalFilename = originalSearchPath? path.relative(originalSearchPath, from) : null;
+        var originalSearchPath = this.getSearchPath(from);
+        var originalFilename = this.getTemplateName(from);
 
         // If we are including same file from a different search path
         // Slice the search paths to avoid including from previous ones