diff options
Diffstat (limited to 'jinja2/optimizer.py')
-rw-r--r-- | jinja2/optimizer.py | 20 |
1 files changed, 16 insertions, 4 deletions
diff --git a/jinja2/optimizer.py b/jinja2/optimizer.py index 5877c6f..4dbd5d9 100644 --- a/jinja2/optimizer.py +++ b/jinja2/optimizer.py @@ -123,6 +123,7 @@ class Optimizer(NodeTransformer): def visit_For(self, node, context): """Loop unrolling for iterable constant values.""" + fallback = self.generic_visit(node.copy(), context) try: iterable = self.visit(node.iter, context).as_const() # we only unroll them if they have a length and are iterable @@ -131,11 +132,8 @@ class Optimizer(NodeTransformer): # we also don't want unrolling if macros are defined in it if node.find(nodes.Macro) is not None: raise TypeError() - # XXX: add support for loop test clauses in the optimizer - if node.test is not None: - raise TypeError() except (nodes.Impossible, TypeError): - return self.generic_visit(node, context) + return fallback parent = context.get('loop') context.push() @@ -157,6 +155,20 @@ class Optimizer(NodeTransformer): else: raise AssertionError('unexpected assignable node') + if node.test is not None: + filtered_sequence = [] + for item in iterable: + context.push() + assign(node.target, item) + try: + rv = self.visit(node.test.copy(), context).as_const() + except: + return fallback + context.pop() + if rv: + filtered_sequence.append(item) + iterable = filtered_sequence + try: try: for item, loop in LoopContext(iterable, parent, True): |