Implement jinja2.ext.loopcontrols extensions

ccutrer · ccutrer · commit 69b70ba81863 · 2024-11-23T12:49:05.000-07:00
`break` and `continue` in for loops. See https://jinja.palletsprojects.com/en/stable/extensions/#loop-controls
diff --git a/src/main/java/com/hubspot/jinjava/interpret/NotInLoopException.java b/src/main/java/com/hubspot/jinjava/interpret/NotInLoopException.java
@@ -0,0 +1,14 @@
+package com.hubspot.jinjava.interpret;
+
+/**
+ * Exception thrown when `continue` or `break` is called outside of a loop
+ */
+public class NotInLoopException extends InterpretException {
+
+  public static final String MESSAGE_PREFIX = "`";
+  public static final String MESSAGE_SUFFIX = "` called while not in a for loop";
+
+  public NotInLoopException(String tagName) {
+    super(MESSAGE_PREFIX + tagName + MESSAGE_SUFFIX);
+  }
+}
diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/BreakTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/BreakTag.java
@@ -0,0 +1,51 @@
+package com.hubspot.jinjava.lib.tag;
+
+import com.hubspot.jinjava.doc.annotations.JinjavaDoc;
+import com.hubspot.jinjava.doc.annotations.JinjavaTextMateSnippet;
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
+import com.hubspot.jinjava.interpret.NotInLoopException;
+import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.util.ForLoop;
+
+/**
+ * Implements the common loopcontrol `continue`, as in the jinja2.ext.loopcontrols extension
+ * @author ccutrer
+ */
+
+@JinjavaDoc(
+  value = "Stops executing the current for loop, including any further iterations"
+)
+@JinjavaTextMateSnippet(
+  code = "{% for item in [1, 2, 3, 4] %}{% if item > 2 == 0 %}{% break %}{% endif %}{{ item }}{% endfor %}"
+)
+public class BreakTag implements Tag {
+
+  public static final String TAG_NAME = "break";
+
+  @Override
+  public String getName() {
+    return TAG_NAME;
+  }
+
+  @Override
+  public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
+    Object loop = interpreter.getContext().get(ForTag.LOOP);
+    if (loop instanceof ForLoop) {
+      ForLoop forLoop = (ForLoop) loop;
+      forLoop.doBreak();
+    } else {
+      throw new NotInLoopException(TAG_NAME);
+    }
+    return "";
+  }
+
+  @Override
+  public String getEndTagName() {
+    return null;
+  }
+
+  @Override
+  public boolean isRenderedInValidationMode() {
+    return true;
+  }
+}
diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/ContinueTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/ContinueTag.java
@@ -0,0 +1,49 @@
+package com.hubspot.jinjava.lib.tag;
+
+import com.hubspot.jinjava.doc.annotations.JinjavaDoc;
+import com.hubspot.jinjava.doc.annotations.JinjavaTextMateSnippet;
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
+import com.hubspot.jinjava.interpret.NotInLoopException;
+import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.util.ForLoop;
+
+/**
+ * Implements the common loopcontrol `continue`, as in the jinja2.ext.loopcontrols extension
+ * @author ccutrer
+ */
+
+@JinjavaDoc(value = "Stops executing the current iteration of the current for loop")
+@JinjavaTextMateSnippet(
+  code = "{% for item in [1, 2, 3, 4] %}{% if item % 2 == 0 %}{% continue %}{% endif %}{{ item }}{% endfor %}"
+)
+public class ContinueTag implements Tag {
+
+  public static final String TAG_NAME = "continue";
+
+  @Override
+  public String getName() {
+    return TAG_NAME;
+  }
+
+  @Override
+  public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
+    Object loop = interpreter.getContext().get(ForTag.LOOP);
+    if (loop instanceof ForLoop) {
+      ForLoop forLoop = (ForLoop) loop;
+      forLoop.doContinue();
+    } else {
+      throw new NotInLoopException(TAG_NAME);
+    }
+    return "";
+  }
+
+  @Override
+  public String getEndTagName() {
+    return null;
+  }
+
+  @Override
+  public boolean isRenderedInValidationMode() {
+    return true;
+  }
+}
diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/ForTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/ForTag.java
@@ -277,6 +277,10 @@ public String renderForCollection(
               interpreter.addError(TemplateError.fromOutputTooBigException(e));
               return checkLoopVariable(interpreter, buff);
             }
+            // continue in the body of the loop; ignore the rest of the body
+            if (loop.isContinued()) {
+              break;
+            }
           }
         }
         if (
@@ -297,7 +301,7 @@ private String checkLoopVariable(
     JinjavaInterpreter interpreter,
     LengthLimitingStringBuilder buff
   ) {
-    if (interpreter.getContext().get("loop") instanceof DeferredValue) {
+    if (interpreter.getContext().get(LOOP) instanceof DeferredValue) {
       throw new DeferredValueException(
         "loop variable deferred",
         interpreter.getLineNumber(),
diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/TagLibrary.java b/src/main/java/com/hubspot/jinjava/lib/tag/TagLibrary.java
@@ -29,7 +29,9 @@ protected void registerDefaults() {
     registerClasses(
       AutoEscapeTag.class,
       BlockTag.class,
+      BreakTag.class,
       CallTag.class,
+      ContinueTag.class,
       CycleTag.class,
       ElseTag.class,
       ElseIfTag.class,
diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerForTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerForTag.java
@@ -179,8 +179,8 @@ private EagerExecutionResult runLoopOnce(
   ) {
     return EagerContextWatcher.executeInChildContext(
       eagerInterpreter -> {
-        if (!(eagerInterpreter.getContext().get("loop") instanceof DeferredValue)) {
-          eagerInterpreter.getContext().put("loop", DeferredValue.instance());
+        if (!(eagerInterpreter.getContext().get(ForTag.LOOP) instanceof DeferredValue)) {
+          eagerInterpreter.getContext().put(ForTag.LOOP, DeferredValue.instance());
         }
         List<String> loopVars = getTag()
           .getLoopVarsAndExpression((TagToken) tagNode.getMaster())
diff --git a/src/main/java/com/hubspot/jinjava/util/EagerContextWatcher.java b/src/main/java/com/hubspot/jinjava/util/EagerContextWatcher.java
@@ -11,6 +11,7 @@
 import com.hubspot.jinjava.interpret.MetaContextVariables;
 import com.hubspot.jinjava.interpret.OneTimeReconstructible;
 import com.hubspot.jinjava.interpret.RevertibleObject;
+import com.hubspot.jinjava.lib.tag.ForTag;
 import com.hubspot.jinjava.lib.tag.eager.EagerExecutionResult;
 import com.hubspot.jinjava.objects.collections.PyList;
 import com.hubspot.jinjava.objects.collections.PyMap;
@@ -211,7 +212,7 @@ private static Map<String, Object> getBasicSpeculativeBindings(
         )
       )
       .filter(entry -> !ignoredKeys.contains(entry.getKey()))
-      .filter(entry -> !"loop".equals(entry.getKey()))
+      .filter(entry -> !ForTag.LOOP.equals(entry.getKey()))
       .map(entry -> {
         if (
           eagerExecutionResult.getResult().isFullyResolved() ||
diff --git a/src/main/java/com/hubspot/jinjava/util/ForLoop.java b/src/main/java/com/hubspot/jinjava/util/ForLoop.java
@@ -28,6 +28,8 @@ public class ForLoop implements Iterator<Object> {
   private int length = NULL_VALUE;
   private boolean first = true;
   private boolean last;
+  private boolean continued;
+  private boolean broken;
 
   private int depth;
 
@@ -45,6 +47,8 @@ public ForLoop(Iterator<?> ite, int len) {
       last = false;
     }
     it = ite;
+    continued = false;
+    broken = false;
   }
 
   public ForLoop(Iterator<?> ite) {
@@ -57,10 +61,16 @@ public ForLoop(Iterator<?> ite) {
       revcounter = 2;
       last = true;
     }
+    continued = false;
+    broken = false;
   }
 
   @Override
   public Object next() {
+    if (broken) {
+      return null;
+    }
+    continued = false;
     Object res;
     if (it.hasNext()) {
       index++;
@@ -129,8 +139,24 @@ public boolean isLast() {
     return last;
   }
 
+  public boolean isContinued() {
+    return continued;
+  }
+
+  public void doContinue() {
+    continued = true;
+  }
+
+  public void doBreak() {
+    continued = true;
+    broken = true;
+  }
+
   @Override
   public boolean hasNext() {
+    if (broken) {
+      return false;
+    }
     return it.hasNext();
   }
 
diff --git a/src/test/java/com/hubspot/jinjava/lib/tag/BreakTagTest.java b/src/test/java/com/hubspot/jinjava/lib/tag/BreakTagTest.java
@@ -0,0 +1,51 @@
+package com.hubspot.jinjava.lib.tag;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.hubspot.jinjava.BaseInterpretingTest;
+import com.hubspot.jinjava.interpret.RenderResult;
+import com.hubspot.jinjava.interpret.TemplateError;
+import org.junit.Test;
+
+public class BreakTagTest extends BaseInterpretingTest {
+
+  @Test
+  public void testBreak() {
+    String template =
+      "{% for item in [1, 2, 3, 4] %}{% if item > 2 %}{% break %}{% endif %}{{ item }}{% endfor %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("12");
+  }
+
+  @Test
+  public void testNestedBreak() {
+    String template =
+      "{% for item in [1, 2, 3, 4] %}{% for item2 in [5, 6, 7] %}{% break %}{{ item2 }}{% endfor %}{{ item }}{% endfor %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("1234");
+  }
+
+  @Test
+  public void testBreakWithEarlierContent() {
+    String template =
+      "{% for item in [1, 2, 3, 4] %}{{ item }}{% if item > 2 %}{% break %}{% endif %}{{ item }}{% endfor %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("11223");
+  }
+
+  @Test
+  public void testBreakOutOfContext() {
+    String template = "{% break %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("");
+    assertThat(rendered.getErrors()).hasSize(1);
+    assertThat(rendered.getErrors().get(0).getSeverity())
+      .isEqualTo(TemplateError.ErrorType.FATAL);
+    assertThat(rendered.getErrors().get(0).getMessage())
+      .contains("NotInLoopException: `break` called while not in a for loop");
+  }
+}
diff --git a/src/test/java/com/hubspot/jinjava/lib/tag/ContinueTagTest.java b/src/test/java/com/hubspot/jinjava/lib/tag/ContinueTagTest.java
@@ -0,0 +1,51 @@
+package com.hubspot.jinjava.lib.tag;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.hubspot.jinjava.BaseInterpretingTest;
+import com.hubspot.jinjava.interpret.RenderResult;
+import com.hubspot.jinjava.interpret.TemplateError;
+import org.junit.Test;
+
+public class ContinueTagTest extends BaseInterpretingTest {
+
+  @Test
+  public void testContinue() {
+    String template =
+      "{% for item in [1, 2, 3, 4] %}{% if item % 2 == 0 %}{% continue %}{% endif %}{{ item }}{% endfor %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("13");
+  }
+
+  @Test
+  public void testNestedContinue() {
+    String template =
+      "{% for item in [1, 2, 3, 4] %}{% for item2 in [5, 6, 7] %}{% continue %}{{ item2 }}{% endfor %}{{ item }}{% endfor %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("1234");
+  }
+
+  @Test
+  public void testContinueWithEarlierContent() {
+    String template =
+      "{% for item in [1, 2, 3, 4] %}{{ item }}{% if item % 2 == 0 %}{% continue %}{% endif %}{{ item }}{% endfor %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("112334");
+  }
+
+  @Test
+  public void testContinueOutOfContext() {
+    String template = "{% continue %}";
+
+    RenderResult rendered = jinjava.renderForResult(template, context);
+    assertThat(rendered.getOutput()).isEqualTo("");
+    assertThat(rendered.getErrors()).hasSize(1);
+    assertThat(rendered.getErrors().get(0).getSeverity())
+      .isEqualTo(TemplateError.ErrorType.FATAL);
+    assertThat(rendered.getErrors().get(0).getMessage())
+      .contains("NotInLoopException: `continue` called while not in a for loop");
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -277,6 +277,10 @@ public String renderForCollection(`
`277`	`277`	`interpreter.addError(TemplateError.fromOutputTooBigException(e));`
`278`	`278`	`return checkLoopVariable(interpreter, buff);`
`279`	`279`	`}`
	`280`	`+ // continue in the body of the loop; ignore the rest of the body`
	`281`	`+ if (loop.isContinued()) {`
	`282`	`+ break;`
	`283`	`+ }`
`280`	`284`	`}`
`281`	`285`	`}`
`282`	`286`	`if (`
`@@ -297,7 +301,7 @@ private String checkLoopVariable(`
`297`	`301`	`JinjavaInterpreter interpreter,`
`298`	`302`	`LengthLimitingStringBuilder buff`
`299`	`303`	`) {`
`300`		`- if (interpreter.getContext().get("loop") instanceof DeferredValue) {`
	`304`	`+ if (interpreter.getContext().get(LOOP) instanceof DeferredValue) {`
`301`	`305`	`throw new DeferredValueException(`
`302`	`306`	`"loop variable deferred",`
`303`	`307`	`interpreter.getLineNumber(),`