Add python slice support to jinja-compat. fixes mozilla#188

CHANGELOG.md

@@ -7,6 +7,8 @@ master (unreleased)
 
 * Fix handling methods and attributes of static arrays, objects and primitives.
   Solves the issue [#937](https://github.com/mozilla/nunjucks/issues/937)
+* Add support for python-style array slices with Jinja compat enabled. Thanks
+  Frankie Dintino.
 
 3.0.0 (Nov 5 2016)
 ----------------

docs/api.md

@@ -135,8 +135,9 @@ nunjucks does not aim for complete Jinja/Python compatibility, this
 might help users seeking just that.
 
 This adds `True` and `False` which map to the JS `true` and `false`
-values, as well as augmenting arrays and objects with Python-style
-methods. [Check out the source](https://github.com/mozilla/nunjucks/blob/master/src/jinja-compat.js)
+values, allows use of Python slice syntax, and augments arrays and
+objects with Python-style methods.
+[Check out the source](https://github.com/mozilla/nunjucks/blob/master/src/jinja-compat.js)
 to see everything it adds.
 {% endapi %}
 {% raw %}

src/jinja-compat.js

@@ -5,8 +5,23 @@ function installCompat() {
     // references the nunjucks instance
     var runtime = this.runtime; // jshint ignore:line
     var lib = this.lib; // jshint ignore:line
+    var Compiler = this.compiler.Compiler; // jshint ignore:line
+    var Parser = this.parser.Parser; // jshint ignore:line
+    var nodes = this.nodes; // jshint ignore:line
+    var lexer = this.lexer; // jshint ignore:line
 
     var orig_contextOrFrameLookup = runtime.contextOrFrameLookup;
+    var orig_Compiler_assertType = Compiler.prototype.assertType;
+    var orig_Parser_parseAggregate = Parser.prototype.parseAggregate;
+    var orig_memberLookup = runtime.memberLookup;
+
+    function uninstall() {
+        runtime.contextOrFrameLookup = orig_contextOrFrameLookup;
+        Compiler.prototype.assertType = orig_Compiler_assertType;
+        Parser.prototype.parseAggregate = orig_Parser_parseAggregate;
+        runtime.memberLookup = orig_memberLookup;
+    }
+
     runtime.contextOrFrameLookup = function(context, frame, key) {
         var val = orig_contextOrFrameLookup.apply(this, arguments);
         if (val === undefined) {
@@ -23,7 +38,132 @@ function installCompat() {
         return val;
     };
 
-    var orig_memberLookup = runtime.memberLookup;
+    var Slice = nodes.Node.extend('Slice', {
+        fields: ['start', 'stop', 'step'],
+        init: function(lineno, colno, start, stop, step) {
+            start = start || new nodes.Literal(lineno, colno, null);
+            stop = stop || new nodes.Literal(lineno, colno, null);
+            step = step || new nodes.Literal(lineno, colno, 1);
+            this.parent(lineno, colno, start, stop, step);
+        }
+    });
+
+    Compiler.prototype.assertType = function(node) {
+        if (node instanceof Slice) {
+            return;
+        }
+        return orig_Compiler_assertType.apply(this, arguments);
+    };
+    Compiler.prototype.compileSlice = function(node, frame) {
+        this.emit('(');
+        this._compileExpression(node.start, frame);
+        this.emit('),(');
+        this._compileExpression(node.stop, frame);
+        this.emit('),(');
+        this._compileExpression(node.step, frame);
+        this.emit(')');
+    };
+
+    function getTokensState(tokens) {
+        return {
+            index: tokens.index,
+            lineno: tokens.lineno,
+            colno: tokens.colno
+        };
+    }
+
+    Parser.prototype.parseAggregate = function() {
+        var self = this;
+        var origState = getTokensState(this.tokens);
+        // Set back one accounting for opening bracket/parens
+        origState.colno--;
+        origState.index--;
+        try {
+            return orig_Parser_parseAggregate.apply(this);
+        } catch(e) {
+            var errState = getTokensState(this.tokens);
+            var rethrow = function() {
+                lib.extend(self.tokens, errState);
+                return e;
+            };
+
+            // Reset to state before original parseAggregate called
+            lib.extend(this.tokens, origState);
+            this.peeked = false;
+
+            var tok = this.peekToken();
+            if (tok.type !== lexer.TOKEN_LEFT_BRACKET) {
+                throw rethrow();
+            } else {
+                this.nextToken();
+            }
+
+            var node = new Slice(tok.lineno, tok.colno);
+
+            // If we don't encounter a colon while parsing, this is not a slice,
+            // so re-raise the original exception.
+            var isSlice = false;
+
+            for (var i = 0; i <= node.fields.length; i++) {
+                if (this.skip(lexer.TOKEN_RIGHT_BRACKET)) {
+                    break;
+                }
+                if (i === node.fields.length) {
+                    if (isSlice) {
+                        this.fail('parseSlice: too many slice components', tok.lineno, tok.colno);
+                    } else {
+                        break;
+                    }
+                }
+                if (this.skip(lexer.TOKEN_COLON)) {
+                    isSlice = true;
+                } else {
+                    var field = node.fields[i];
+                    node[field] = this.parseExpression();
+                    isSlice = this.skip(lexer.TOKEN_COLON) || isSlice;
+                }
+            }
+            if (!isSlice) {
+                throw rethrow();
+            }
+            return new nodes.Array(tok.lineno, tok.colno, [node]);
+        }
+    };
+
+    function sliceLookup(obj, start, stop, step) {
+        obj = obj || [];
+        if (start === null) {
+            start = (step < 0) ? (obj.length - 1) : 0;
+        }
+        if (stop === null) {
+            stop = (step < 0) ? -1 : obj.length;
+        } else {
+            if (stop < 0) {
+                stop += obj.length;
+            }
+        }
+
+        if (start < 0) {
+            start += obj.length;
+        }
+
+        var results = [];
+
+        for (var i = start; ; i += step) {
+            if (i < 0 || i > obj.length) {
+                break;
+            }
+            if (step > 0 && i >= stop) {
+                break;
+            }
+            if (step < 0 && i <= stop) {
+                break;
+            }
+            results.push(runtime.memberLookup(obj, i));
+        }
+        return results;
+    }
+
     var ARRAY_MEMBERS = {
         pop: function(index) {
             if (index === undefined) {
@@ -140,7 +280,10 @@ function installCompat() {
     OBJECT_MEMBERS.itervalues = OBJECT_MEMBERS.values;
     OBJECT_MEMBERS.iterkeys = OBJECT_MEMBERS.keys;
     runtime.memberLookup = function(obj, val, autoescape) { // jshint ignore:line
-        obj = obj || {};
+        if (arguments.length === 4) {
+            return sliceLookup.apply(this, arguments);
+        }
+        obj = obj || {};    
 
         // If the object is an object, return any of the methods that Python would
         // otherwise provide.
@@ -154,6 +297,8 @@ function installCompat() {
 
         return orig_memberLookup.apply(this, arguments);
     };
+
+    return uninstall;
 }
 
 module.exports = installCompat;

src/parser.js

@@ -1294,5 +1294,6 @@ module.exports = {
             p.extensions = extensions;
         }
         return p.parseAsRoot();
-    }
+    },
+    Parser: Parser
 };

tests/jinja-compat.js

@@ -0,0 +1,100 @@
+(function() {
+    'use strict';
+
+    var nunjucks, util;
+
+    if(typeof require !== 'undefined') {
+        nunjucks = require('../index.js');
+        util = require('./util');
+    }
+    else {
+        nunjucks = window.nunjucks;
+        util = window.util;
+    }
+
+    var equal = util.jinjaEqual;
+    var finish = util.finish;
+
+    describe('jinja-compat', function() {
+        var arr = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'];
+
+        it('should support array slices with start and stop', function(done) {
+            equal('{% for i in arr[1:4] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'bcd');
+            finish(done);
+        });
+        it('should support array slices using expressions', function(done) {
+            equal('{% for i in arr[n:n+3] %}{{ i }}{% endfor %}',
+                  { n: 1, arr: arr },
+                  'bcd');
+            finish(done);
+        });
+        it('should support array slices with start', function(done) {
+            equal('{% for i in arr[3:] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'defgh');
+            finish(done);
+        });
+        it('should support array slices with negative start', function(done) {
+            equal('{% for i in arr[-3:] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'fgh');
+            finish(done);
+        });
+        it('should support array slices with stop', function(done) {
+            equal('{% for i in arr[:4] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'abcd');
+            finish(done);
+        });
+        it('should support array slices with negative stop', function(done) {
+            equal('{% for i in arr[:-3] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'abcde');
+            finish(done);
+        });
+        it('should support array slices with step', function(done) {
+            equal('{% for i in arr[::2] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'aceg');
+            finish(done);
+        });
+        it('should support array slices with negative step', function(done) {
+            equal('{% for i in arr[::-1] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'hgfedcba');
+            finish(done);
+        });
+        it('should support array slices with start and negative step', function(done) {
+            equal('{% for i in arr[4::-1] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'edcba');
+            finish(done);
+        });
+        it('should support array slices with negative start and negative step', function(done) {
+            equal('{% for i in arr[-5::-1] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'dcba');
+            finish(done);
+        });
+        it('should support array slices with stop and negative step', function(done) {
+            equal('{% for i in arr[:3:-1] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'hgfe');
+            finish(done);
+        });
+        it('should support array slices with start and step', function(done) {
+            equal('{% for i in arr[1::2] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'bdfh');
+            finish(done);
+        });
+        it('should support array slices with start, stop, and step', function(done) {
+            equal('{% for i in arr[1:7:2] %}{{ i }}{% endfor %}',
+                  { arr: arr },
+                  'bdf');
+            finish(done);
+        });
+    });
+})();

tests/util.js

@@ -1,22 +1,21 @@
 (function() {
     'use strict';
 
-    var Environment, Template, Loader, templatesPath, expect;
+    var nunjucks, Environment, Template, Loader, templatesPath, expect;
 
     if(typeof require !== 'undefined') {
-        Environment = require('../src/environment').Environment;
-        Template = require('../src/environment').Template;
-        Loader = require('../src/node-loaders').FileSystemLoader;
+        nunjucks = require('../index.js');
+        Loader = nunjucks.FileSystemLoader;
         templatesPath = 'tests/templates';
         expect = require('expect.js');
     }
     else {
-        Environment = nunjucks.Environment;
-        Template = nunjucks.Template;
         Loader = nunjucks.WebLoader;
         templatesPath = '../templates';
         expect = window.expect;
     }
+    Environment = nunjucks.Environment;
+    Template = nunjucks.Template;
 
     var numAsyncs;
     var doneHandler;
@@ -37,6 +36,15 @@
         expect(res).to.be(str2);
     }
 
+    function jinjaEqual(str, ctx, str2, env) {
+        var jinjaUninstall = nunjucks.installJinjaCompat();
+        try {
+            return equal(str, ctx, str2, env);
+        } finally {
+            jinjaUninstall();
+        }
+    }
+
     function finish(done) {
         if(numAsyncs > 0) {
             doneHandler = done;
@@ -118,13 +126,15 @@
     if(typeof module !== 'undefined') {
         module.exports.render = render;
         module.exports.equal = equal;
+        module.exports.jinjaEqual = jinjaEqual;
         module.exports.finish = finish;
         module.exports.normEOL = normEOL;
     }
     else {
         window.util = {
             render: render,
             equal: equal,
+            jinjaEqual: jinjaEqual,
             finish: finish,
             normEOL: normEOL
         };