optimizer can optimize filtered for loops now

--HG--
branch : trunk

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):