Add DynamicRenderedOutputNode to allow path reconstruction when resolving blocks

jasmith-hs · jasmith-hs · commit 7cb442ebeb44 · 2024-05-06T11:07:12.000-04:00
diff --git a/src/main/java/com/hubspot/jinjava/interpret/JinjavaInterpreter.java b/src/main/java/com/hubspot/jinjava/interpret/JinjavaInterpreter.java
@@ -36,6 +36,7 @@
 import com.hubspot.jinjava.lib.tag.DoTag;
 import com.hubspot.jinjava.lib.tag.ExtendsTag;
 import com.hubspot.jinjava.lib.tag.eager.EagerGenericTag;
+import com.hubspot.jinjava.lib.tag.eager.importing.EagerImportingStrategyFactory;
 import com.hubspot.jinjava.objects.serialization.PyishObjectMapper;
 import com.hubspot.jinjava.objects.serialization.PyishSerializable;
 import com.hubspot.jinjava.random.ConstantZeroRandomNumberGenerator;
@@ -46,6 +47,7 @@
 import com.hubspot.jinjava.tree.TreeParser;
 import com.hubspot.jinjava.tree.output.BlockInfo;
 import com.hubspot.jinjava.tree.output.BlockPlaceholderOutputNode;
+import com.hubspot.jinjava.tree.output.DynamicRenderedOutputNode;
 import com.hubspot.jinjava.tree.output.OutputList;
 import com.hubspot.jinjava.tree.output.OutputNode;
 import com.hubspot.jinjava.tree.output.RenderedOutputNode;
@@ -513,7 +515,10 @@ private void resolveBlockStubs(OutputList output, Stack<String> blockNames) {
           currentBlock = block;
 
           OutputList blockValueBuilder = new OutputList(config.getMaxOutputSize());
+          DynamicRenderedOutputNode prefix = new DynamicRenderedOutputNode();
+          blockValueBuilder.addNode(prefix);
           boolean pushedParentPathOntoStack = false;
+          int numDeferredTokensBefore = context.getDeferredTokens().size();
           if (
             block.getParentPath().isPresent() &&
             !getContext().getCurrentPathStack().contains(block.getParentPath().get())
@@ -535,7 +540,19 @@ private void resolveBlockStubs(OutputList output, Stack<String> blockNames) {
             blockValueBuilder.addNode(child.render(this));
           }
           if (pushedParentPathOntoStack) {
-            getContext().getCurrentPathStack().pop();
+            if (context.getDeferredTokens().size() > numDeferredTokensBefore) {
+              prefix.setValue(
+                EagerImportingStrategyFactory.getSetTagForCurrentPath(this)
+              );
+              getContext().getCurrentPathStack().pop();
+              blockValueBuilder.addNode(
+                new RenderedOutputNode(
+                  EagerImportingStrategyFactory.getSetTagForCurrentPath(this)
+                )
+              );
+            } else {
+              getContext().getCurrentPathStack().pop();
+            }
           }
           blockNames.push(blockPlaceholder.getBlockName());
           resolveBlockStubs(blockValueBuilder, blockNames);
diff --git a/src/main/java/com/hubspot/jinjava/tree/output/DynamicRenderedOutputNode.java b/src/main/java/com/hubspot/jinjava/tree/output/DynamicRenderedOutputNode.java
@@ -0,0 +1,33 @@
+package com.hubspot.jinjava.tree.output;
+
+import com.google.common.base.Charsets;
+import java.nio.charset.Charset;
+
+/**
+ * An OutputNode that can be modified after already being added to the OutputList
+ */
+public class DynamicRenderedOutputNode implements OutputNode {
+
+  protected String output = "";
+
+  public void setValue(String output) {
+    this.output = output;
+  }
+
+  @Override
+  public String getValue() {
+    return output;
+  }
+
+  @Override
+  public long getSize() {
+    return output == null
+      ? 0
+      : output.getBytes(Charset.forName(Charsets.UTF_8.name())).length;
+  }
+
+  @Override
+  public String toString() {
+    return getValue();
+  }
+}
