Keep meta context variables when importing

src/main/java/com/hubspot/jinjava/el/ext/eager/EvalResultHolder.java

@@ -136,7 +136,7 @@ static String reconstructNode(
                 .getELResolver()
                 .getValue(context, null, ExtendedParser.INTERPRETER)
             ).getContext()
-              .getMetaContextVariables()
+              .getComputedMetaContextVariables()
               .contains(name)
           ) {
             return name;

src/main/java/com/hubspot/jinjava/interpret/Context.java

@@ -20,6 +20,7 @@
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.SetMultimap;
+import com.google.common.collect.Sets;
 import com.hubspot.jinjava.lib.Importable;
 import com.hubspot.jinjava.lib.expression.DefaultExpressionStrategy;
 import com.hubspot.jinjava.lib.expression.ExpressionStrategy;
@@ -34,6 +35,7 @@
 import com.hubspot.jinjava.lib.tag.Tag;
 import com.hubspot.jinjava.lib.tag.TagLibrary;
 import com.hubspot.jinjava.lib.tag.eager.DeferredToken;
+import com.hubspot.jinjava.mode.EagerExecutionMode;
 import com.hubspot.jinjava.tree.Node;
 import com.hubspot.jinjava.util.DeferredValueUtils;
 import com.hubspot.jinjava.util.ScopeMap;
@@ -117,6 +119,7 @@ public enum Library {
   private boolean unwrapRawOverride = false;
   private DynamicVariableResolver dynamicVariableResolver = null;
   private final Set<String> metaContextVariables; // These variable names aren't tracked in eager execution
+  private final Set<String> overriddenNonMetaContextVariables;
   private Node currentNode;
 
   public Context() {
@@ -209,6 +212,8 @@ public Context(
       new FunctionLibrary(parent == null, disabled.get(Library.FUNCTION));
     this.metaContextVariables =
       parent == null ? new HashSet<>() : parent.metaContextVariables;
+    this.overriddenNonMetaContextVariables =
+      parent == null ? new HashSet<>() : parent.overriddenNonMetaContextVariables;
     if (parent != null) {
       this.expressionStrategy = parent.expressionStrategy;
       this.partialMacroEvaluation = parent.partialMacroEvaluation;
@@ -348,10 +353,37 @@ public void addResolvedFunction(String function) {
     }
   }
 
+  @Deprecated
+  @Beta
   public Set<String> getMetaContextVariables() {
     return metaContextVariables;
   }
 
+  @Beta
+  public Set<String> getComputedMetaContextVariables() {
+    return Sets.difference(metaContextVariables, overriddenNonMetaContextVariables);
+  }
+
+  @Beta
+  public void addMetaContextVariables(Collection<String> variables) {
+    metaContextVariables.addAll(variables);
+  }
+
+  @Beta
+  public void addNonMetaContextVariables(Collection<String> variables) {
+    overriddenNonMetaContextVariables.addAll(
+      variables
+        .stream()
+        .filter(var -> !EagerExecutionMode.STATIC_META_CONTEXT_VARIABLES.contains(var))
+        .collect(Collectors.toList())
+    );
+  }
+
+  @Beta
+  public void removeNonMetaContextVariables(Collection<String> variables) {
+    overriddenNonMetaContextVariables.removeAll(variables);
+  }
+
   public void handleDeferredNode(Node node) {
     if (
       JinjavaInterpreter

src/main/java/com/hubspot/jinjava/interpret/JinjavaInterpreter.java

@@ -111,7 +111,6 @@ public JinjavaInterpreter(
     this.context = context;
     this.config = renderConfig;
     this.application = application;
-
     this.config.getExecutionMode().prepareContext(this.context);
 
     switch (config.getRandomNumberGeneratorStrategy()) {

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

@@ -37,7 +37,6 @@
 import com.hubspot.jinjava.tree.Node;
 import com.hubspot.jinjava.tree.TagNode;
 import com.hubspot.jinjava.tree.parse.TagToken;
-import com.hubspot.jinjava.util.EagerReconstructionUtils;
 import com.hubspot.jinjava.util.ForLoop;
 import com.hubspot.jinjava.util.HelperStringTokenizer;
 import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
@@ -48,7 +47,6 @@
 import java.util.List;
 import java.util.Map.Entry;
 import java.util.Optional;
-import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import org.apache.commons.lang3.tuple.Pair;
@@ -165,11 +163,7 @@ public String renderForCollection(
   ) {
     ForLoop loop = ObjectIterator.getLoop(collection);
 
-    Set<String> removedMetaContextVariables =
-      EagerReconstructionUtils.removeMetaContextVariables(
-        loopVars.stream(),
-        interpreter.getContext()
-      );
+    interpreter.getContext().addNonMetaContextVariables(loopVars);
     try (InterpreterScopeClosable c = interpreter.enterScope()) {
       if (interpreter.isValidationMode() && !loop.hasNext()) {
         loop = ObjectIterator.getLoop(new DummyObject());
@@ -298,10 +292,7 @@ public String renderForCollection(
       }
       return checkLoopVariable(interpreter, buff);
     } finally {
-      interpreter
-        .getContext()
-        .getMetaContextVariables()
-        .addAll(removedMetaContextVariables);
+      interpreter.getContext().removeNonMetaContextVariables(loopVars);
     }
   }
 

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

@@ -378,7 +378,7 @@ private static Collection<String> markDeferredWordsAndFindSources(
         }
         return !(val instanceof DeferredValue);
       })
-      .filter(prop -> !context.getMetaContextVariables().contains(prop))
+      .filter(prop -> !context.getComputedMetaContextVariables().contains(prop))
       .filter(prop -> {
         DeferredValue deferredValue = convertToDeferredValue(context, prop);
         context.put(prop, deferredValue);

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

@@ -185,11 +185,7 @@ private EagerExecutionResult runLoopOnce(
         List<String> loopVars = getTag()
           .getLoopVarsAndExpression((TagToken) tagNode.getMaster())
           .getLeft();
-        Set<String> removedMetaContextVariables =
-          EagerReconstructionUtils.removeMetaContextVariables(
-            loopVars.stream(),
-            interpreter.getContext()
-          );
+        interpreter.getContext().addNonMetaContextVariables(loopVars);
         loopVars.forEach(var ->
           interpreter.getContext().put(var, DeferredValue.instance())
         );
@@ -198,10 +194,7 @@ private EagerExecutionResult runLoopOnce(
             renderChildren(tagNode, eagerInterpreter)
           );
         } finally {
-          interpreter
-            .getContext()
-            .getMetaContextVariables()
-            .addAll(removedMetaContextVariables);
+          interpreter.getContext().removeNonMetaContextVariables(loopVars);
           if (clearDeferredWords) {
             interpreter
               .getContext()

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

@@ -40,11 +40,11 @@ public String run(TagNode tagNode, JinjavaInterpreter interpreter) {
       variables = new String[] { var };
       expression = tagNode.getHelpers();
     }
-
-    EagerReconstructionUtils.removeMetaContextVariables(
-      Arrays.stream(variables).map(String::trim),
-      interpreter.getContext()
-    );
+    interpreter
+      .getContext()
+      .addNonMetaContextVariables(
+        Arrays.stream(variables).map(String::trim).collect(Collectors.toList())
+      );
 
     EagerExecutionResult eagerExecutionResult = getEagerExecutionResult(
       tagNode,
@@ -167,7 +167,7 @@ public static String getSuffixToPreserveState(
     if (
       maybeTemporaryImportAlias.isPresent() &&
       !AliasedEagerImportingStrategy.isTemporaryImportAlias(variables) &&
-      !interpreter.getContext().getMetaContextVariables().contains(variables)
+      !interpreter.getContext().getComputedMetaContextVariables().contains(variables)
     ) {
       if (!interpreter.getContext().containsKey(maybeTemporaryImportAlias.get())) {
         if (

src/main/java/com/hubspot/jinjava/lib/tag/eager/importing/FlatEagerImportingStrategy.java

@@ -72,7 +72,7 @@ public String getFinalOutput(
       Set<String> metaContextVariables = importingData
         .getOriginalInterpreter()
         .getContext()
-        .getMetaContextVariables();
+        .getComputedMetaContextVariables();
       // defer imported variables
       EagerReconstructionUtils.buildSetTag(
         child

src/main/java/com/hubspot/jinjava/mode/EagerExecutionMode.java

@@ -53,6 +53,6 @@ public void prepareContext(Context context) {
       .filter(Optional::isPresent)
       .forEach(maybeEagerTag -> context.registerTag(maybeEagerTag.get()));
     context.setExpressionStrategy(new EagerExpressionStrategy());
-    context.getMetaContextVariables().addAll(STATIC_META_CONTEXT_VARIABLES);
+    context.addMetaContextVariables(STATIC_META_CONTEXT_VARIABLES);
   }
 }

src/main/java/com/hubspot/jinjava/util/DeferredValueUtils.java

@@ -119,7 +119,7 @@ private static void markDeferredProperties(Context context, Set<String> props) {
     props
       .stream()
       .filter(prop -> !(context.get(prop) instanceof DeferredValue))
-      .filter(prop -> !context.getMetaContextVariables().contains(prop))
+      .filter(prop -> !context.getComputedMetaContextVariables().contains(prop))
       .forEach(prop -> {
         Object value = context.get(prop);
         if (value != null) {

src/main/java/com/hubspot/jinjava/util/EagerContextWatcher.java

@@ -50,7 +50,7 @@ public static EagerExecutionResult executeInChildContext(
   ) {
     final Set<String> metaContextVariables = interpreter
       .getContext()
-      .getMetaContextVariables();
+      .getComputedMetaContextVariables();
     final EagerExecutionResult initialResult;
     final Map<String, Object> speculativeBindings;
     if (eagerChildContextConfig.checkForContextChanges) {

src/main/java/com/hubspot/jinjava/util/EagerReconstructionUtils.java

@@ -2,7 +2,6 @@
 
 import com.google.common.annotations.Beta;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Sets;
 import com.hubspot.jinjava.el.ext.AbstractCallableMethod;
 import com.hubspot.jinjava.interpret.Context;
 import com.hubspot.jinjava.interpret.Context.Library;
@@ -25,7 +24,6 @@
 import com.hubspot.jinjava.lib.tag.eager.importing.AliasedEagerImportingStrategy;
 import com.hubspot.jinjava.lib.tag.eager.importing.EagerImportingStrategyFactory;
 import com.hubspot.jinjava.loader.RelativePathResolver;
-import com.hubspot.jinjava.mode.EagerExecutionMode;
 import com.hubspot.jinjava.objects.serialization.PyishBlockSetSerializable;
 import com.hubspot.jinjava.objects.serialization.PyishObjectMapper;
 import com.hubspot.jinjava.objects.serialization.PyishSerializable;
@@ -313,7 +311,9 @@ private static PrefixToPreserveState hydrateReconstructionOfVariablesBeforeDefer
     JinjavaInterpreter interpreter,
     int depth
   ) {
-    Set<String> metaContextVariables = interpreter.getContext().getMetaContextVariables();
+    Set<String> metaContextVariables = interpreter
+      .getContext()
+      .getComputedMetaContextVariables();
     deferredWords
       .stream()
       .filter(w -> !metaContextVariables.contains(w))
@@ -753,22 +753,6 @@ public static String wrapInChildScope(String toWrap, JinjavaInterpreter interpre
     );
   }
 
-  public static Set<String> removeMetaContextVariables(
-    Stream<String> varStream,
-    Context context
-  ) {
-    Set<String> metaSetVars = Sets
-      .intersection(
-        context.getMetaContextVariables(),
-        varStream
-          .filter(var -> !EagerExecutionMode.STATIC_META_CONTEXT_VARIABLES.contains(var))
-          .collect(Collectors.toSet())
-      )
-      .immutableCopy();
-    context.getMetaContextVariables().removeAll(metaSetVars);
-    return metaSetVars;
-  }
-
   public static Boolean isDeferredExecutionMode() {
     return JinjavaInterpreter
       .getCurrentMaybe()

src/test/java/com/hubspot/jinjava/EagerTest.java

@@ -22,6 +22,7 @@
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
@@ -990,7 +991,7 @@ public void itHandlesImportInDeferredIf() {
 
   @Test
   public void itAllowsMetaContextVarOverriding() {
-    interpreter.getContext().getMetaContextVariables().add("meta");
+    interpreter.getContext().addMetaContextVariables(Collections.singleton("meta"));
     interpreter.getContext().put("meta", "META");
     expectedTemplateInterpreter.assertExpectedOutputNonIdempotent(
       "allows-meta-context-var-overriding"
@@ -1471,7 +1472,7 @@ public void itReconstructsNestedValueInStringRepresentationSecondPass() {
 
   @Test
   public void itDefersLoopSettingMetaContextVar() {
-    interpreter.getContext().getMetaContextVariables().add("content");
+    interpreter.getContext().addMetaContextVariables(Collections.singleton("content"));
     expectedTemplateInterpreter.assertExpectedOutput(
       "defers-loop-setting-meta-context-var"
     );
@@ -1480,7 +1481,7 @@ public void itDefersLoopSettingMetaContextVar() {
   @Test
   public void itDefersLoopSettingMetaContextVarSecondPass() {
     interpreter.getContext().put("deferred", "resolved");
-    interpreter.getContext().getMetaContextVariables().add("content");
+    interpreter.getContext().addMetaContextVariables(Collections.singleton("content"));
     expectedTemplateInterpreter.assertExpectedOutput(
       "defers-loop-setting-meta-context-var.expected"
     );
@@ -1572,4 +1573,21 @@ public void itReconstructsBlockPathWhenDeferredNestedSecondPass() {
       "reconstructs-block-path-when-deferred-nested/test.expected"
     );
   }
+
+  @Test
+  public void itKeepsMetaContextVariablesThroughImport() {
+    setupWithExecutionMode(
+      new EagerExecutionMode() {
+        @Override
+        public void prepareContext(Context context) {
+          super.prepareContext(context);
+          context.addMetaContextVariables(Collections.singleton("meta"));
+        }
+      }
+    );
+    interpreter.getContext().put("meta", new ArrayList<>());
+    expectedTemplateInterpreter.assertExpectedOutput(
+      "keeps-meta-context-variables-through-import/test"
+    );
+  }
 }

src/test/java/com/hubspot/jinjava/el/ext/eager/EagerAstMethodTest.java

@@ -16,6 +16,7 @@
 import com.hubspot.jinjava.objects.collections.PyMap;
 import com.hubspot.jinjava.random.RandomNumberGeneratorStrategy;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -225,7 +226,7 @@ public void itPreservesAstTuple() {
   @Test
   public void itPreservesUnresolvable() {
     interpreter.getContext().put("foo_object", new Foo());
-    interpreter.getContext().getMetaContextVariables().add("foo_object");
+    interpreter.getContext().addMetaContextVariables(Collections.singleton("foo_object"));
     try {
       interpreter.resolveELExpression("foo_object.deferred|upper", -1);
       fail("Should throw DeferredParsingException");

src/test/java/com/hubspot/jinjava/util/EagerExpressionResolverTest.java

@@ -600,7 +600,7 @@ public void itHandlesPyishSerializable() {
   @Test
   public void itHandlesPyishSerializableWithProcessingException() {
     context.put("foo", new SomethingExceptionallyPyish("yes"));
-    context.getMetaContextVariables().add("foo");
+    context.addMetaContextVariables(Collections.singleton("foo"));
     assertThat(interpreter.render("{{ deferred && (1 == 2 || foo) }}"))
       .isEqualTo("{{ deferred && (false || foo) }}");
   }