Handle spaces better in for loop expressions

src/main/java/com/hubspot/jinjava/lib/tag/ForTag.java

@@ -34,6 +34,7 @@
 import com.hubspot.jinjava.tree.ExpressionNode;
 import com.hubspot.jinjava.tree.Node;
 import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.tree.parse.TagToken;
 import com.hubspot.jinjava.util.ForLoop;
 import com.hubspot.jinjava.util.HelperStringTokenizer;
 import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
@@ -43,7 +44,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
-import org.apache.commons.lang3.StringUtils;
+import java.util.Optional;
+import org.apache.commons.lang3.tuple.Pair;
 
 /**
  * {% for a in b|f1:d,c %}
@@ -126,24 +128,16 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
   }
 
   public String interpretUnchecked(TagNode tagNode, JinjavaInterpreter interpreter) {
-    String helpers = getWhitespaceAdjustedHelpers(tagNode.getHelpers());
-    List<String> helper = new HelperStringTokenizer(helpers).splitComma(true).allTokens();
-
-    List<String> loopVars = getLoopVars(helper);
-
-    if (loopVars.size() >= helper.size()) {
-      throw new TemplateSyntaxException(
-        tagNode.getHelpers().trim(),
-        "Tag 'for' expects valid 'in' clause, got: " + tagNode.getHelpers(),
-        tagNode.getLineNumber(),
-        tagNode.getStartPosition()
-      );
-    }
+    Pair<List<String>, String> loopVarsAndExpression = getLoopVarsAndExpression(
+      (TagToken) tagNode.getMaster()
+    );
+    List<String> loopVars = loopVarsAndExpression.getLeft();
+    String loopExpression = loopVarsAndExpression.getRight();
 
-    String loopExpr = getLoopExpression(helper, loopVars);
     Object collection;
     try {
-      collection = interpreter.resolveELExpression(loopExpr, tagNode.getLineNumber());
+      collection =
+        interpreter.resolveELExpression(loopExpression, tagNode.getLineNumber());
     } catch (DeferredParsingException e) {
       throw new DeferredParsingException(
         String.format("%s in %s", String.join(", ", loopVars), e.getDeferredEvalResult())
@@ -233,35 +227,33 @@ public String interpretUnchecked(TagNode tagNode, JinjavaInterpreter interpreter
     }
   }
 
-  public static String getWhitespaceAdjustedHelpers(String helpers) {
-    /* apdlv72@gmail.com
-     * Fix for issues with for-loops that contain whitespace in their range, e.g.
-     * "{% for i in range(1 * 1, 2 * 2) %}"
-     * This is because HelperStringTokenizer will split the range expressions also
-     * at white spaces and end up with [i, in, range(1, *, 1, 2, *, 2)].
-     * To avoid this, the below fix will remove white space from the expression
-     * on the right side of the keyword "in". It will do so however only if there
-     * are no characters in this expression that indicate strings - namely ' and ".
-     * This avoids messing up expressions like {% for i in ['a ','b'] %} that
-     * contain spaces in the arguments.
-     * TODO A somewhat more sophisticated tokenizing/parsing of the for-loop expression.
-     */
-    String[] parts = helpers.split("\\s+in\\s+");
-    if (parts.length == 2 && !parts[1].contains("'") && !parts[1].contains("\"")) {
-      helpers = parts[0] + " in " + parts[1].replace(" ", "");
+  public Pair<List<String>, String> getLoopVarsAndExpression(TagToken tagToken) {
+    List<String> helperTokens = new HelperStringTokenizer(tagToken.getHelpers())
+      .splitComma(true)
+      .allTokens();
+    List<String> loopVars = getLoopVars(helperTokens);
+    Optional<String> maybeLoopExpr = getLoopExpression(tagToken.getHelpers());
+
+    if (loopVars.size() >= helperTokens.size() || !maybeLoopExpr.isPresent()) {
+      throw new TemplateSyntaxException(
+        tagToken.getHelpers().trim(),
+        "Tag 'for' expects valid 'in' clause, got: " + tagToken.getHelpers(),
+        tagToken.getLineNumber(),
+        tagToken.getStartPosition()
+      );
     }
-    return helpers;
+    return Pair.of(loopVars, maybeLoopExpr.get());
   }
 
-  public String getLoopExpression(List<String> helper, List<String> loopVars) {
-    String loopExpr = StringUtils.join(
-      helper.subList(loopVars.size() + 1, helper.size()),
-      ","
-    );
-    return loopExpr;
+  private Optional<String> getLoopExpression(String helpers) {
+    int inIndex = helpers.indexOf(" in ");
+    if (inIndex > 0) {
+      return Optional.of(helpers.substring(inIndex + 4).trim());
+    }
+    return Optional.empty();
   }
 
-  public List<String> getLoopVars(List<String> helper) {
+  private List<String> getLoopVars(List<String> helper) {
     List<String> loopVars = Lists.newArrayList();
     while (loopVars.size() < helper.size()) {
       String val = helper.get(loopVars.size());

src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerForTag.java

@@ -4,18 +4,17 @@
 import com.hubspot.jinjava.interpret.DeferredValueException;
 import com.hubspot.jinjava.interpret.InterpretException;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
-import com.hubspot.jinjava.interpret.TemplateSyntaxException;
 import com.hubspot.jinjava.lib.tag.ForTag;
 import com.hubspot.jinjava.tree.TagNode;
 import com.hubspot.jinjava.tree.parse.TagToken;
 import com.hubspot.jinjava.util.EagerExpressionResolver;
 import com.hubspot.jinjava.util.EagerExpressionResolver.EagerExpressionResult;
-import com.hubspot.jinjava.util.HelperStringTokenizer;
 import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 import com.hubspot.jinjava.util.LengthLimitingStringJoiner;
 import java.util.HashSet;
 import java.util.List;
 import java.util.stream.Collectors;
+import org.apache.commons.lang3.tuple.Pair;
 
 public class EagerForTag extends EagerTagDecorator<ForTag> {
 
@@ -82,22 +81,10 @@ public String eagerInterpret(
 
   @Override
   public String getEagerTagImage(TagToken tagToken, JinjavaInterpreter interpreter) {
-    List<String> helperTokens = new HelperStringTokenizer(
-      ForTag.getWhitespaceAdjustedHelpers(tagToken.getHelpers())
-    )
-      .splitComma(true)
-      .allTokens();
-    List<String> loopVars = getTag().getLoopVars(helperTokens);
-    if (loopVars.size() >= helperTokens.size()) {
-      throw new TemplateSyntaxException(
-        tagToken.getHelpers().trim(),
-        "Tag 'for' expects valid 'in' clause, got: " + tagToken.getHelpers(),
-        tagToken.getLineNumber(),
-        tagToken.getStartPosition()
-      );
-    }
-
-    String loopExpression = getTag().getLoopExpression(helperTokens, loopVars);
+    Pair<List<String>, String> loopVarsAndExpression = getTag()
+      .getLoopVarsAndExpression(tagToken);
+    List<String> loopVars = loopVarsAndExpression.getLeft();
+    String loopExpression = loopVarsAndExpression.getRight();
 
     EagerExpressionResult eagerExpressionResult = EagerExpressionResolver.resolveExpression(
       loopExpression,

src/test/java/com/hubspot/jinjava/lib/tag/ForTagTest.java

@@ -279,6 +279,17 @@ public void forLoopShouldCountUsingNamespaceVariable() {
     assertThat(rendered).isEqualTo("Found item having something: 4");
   }
 
+  @Test
+  public void itShouldHandleSpacesInMaps() {
+    String template =
+      "{% for item in [{'key': 'foo?'}, {'key': 'bar?'}] %}" +
+      "{{ item.key }}\n" +
+      "{% endfor %}";
+
+    String rendered = jinjava.render(template, context);
+    assertThat(rendered).isEqualTo("foo?\nbar?\n");
+  }
+
   private Node fixture(String name) {
     try {
       return new TreeParser(