Add new key that's used to handle deferred imported macros

jasmith-hs · jasmith-hs · commit 0eb7b0c223ec · 2021-03-11T14:35:37.000-05:00
diff --git a/src/main/java/com/hubspot/jinjava/interpret/Context.java b/src/main/java/com/hubspot/jinjava/interpret/Context.java
@@ -51,6 +51,9 @@
 public class Context extends ScopeMap<String, Object> {
   public static final String GLOBAL_MACROS_SCOPE_KEY = "__macros__";
   public static final String IMPORT_RESOURCE_PATH_KEY = "import_resource_path";
+  public static final String DEFERRED_IMPORT_RESOURCE_PATH_KEY =
+    "deferred_import_resource_path";
+
   public static final String IMPORT_RESOURCE_ALIAS_KEY = "import_resource_alias";
 
   private SetMultimap<String, String> dependencies = HashMultimap.create();
diff --git a/src/main/java/com/hubspot/jinjava/lib/fn/eager/EagerMacroFunction.java b/src/main/java/com/hubspot/jinjava/lib/fn/eager/EagerMacroFunction.java
@@ -10,9 +10,7 @@
 import com.hubspot.jinjava.lib.fn.MacroFunction;
 import com.hubspot.jinjava.lib.tag.MacroTag;
 import com.hubspot.jinjava.lib.tag.eager.EagerTagDecorator;
-import com.hubspot.jinjava.objects.collections.PyMap;
 import com.hubspot.jinjava.objects.serialization.PyishObjectMapper;
-import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -110,39 +108,33 @@ public String reconstructImage() {
     Optional<String> importFile = macroFunction.getImportFile(interpreter);
     if (importFile.isPresent()) {
       interpreter.getContext().getCurrentPathStack().pop();
-      if (fullName.indexOf('.') >= 0) {
-        prefix = getAliasedImportResourcePrefix(importFile.get());
-      } else {
-        prefix =
-          EagerTagDecorator.buildSetTagForDeferredInChildContext(
-            ImmutableMap.of(
-              Context.IMPORT_RESOURCE_PATH_KEY,
-              interpreter
-                .getContext()
-                .getPyishObjectMapper()
-                .getAsPyishString(importFile.get())
-            ),
-            interpreter,
-            false
-          );
-        String currentImportResource = interpreter
-          .getContext()
-          .getCurrentPathStack()
-          .peek()
-          .orElse("");
-        suffix =
-          EagerTagDecorator.buildSetTagForDeferredInChildContext(
-            ImmutableMap.of(
-              Context.IMPORT_RESOURCE_PATH_KEY,
-              interpreter
-                .getContext()
-                .getPyishObjectMapper()
-                .getAsPyishString(currentImportResource)
-            ),
-            interpreter,
-            false
-          );
-      }
+      String currentDeferredImportResource = (String) interpreter
+        .getContext()
+        .get(Context.DEFERRED_IMPORT_RESOURCE_PATH_KEY);
+      prefix =
+        EagerTagDecorator.buildSetTagForDeferredInChildContext(
+          ImmutableMap.of(
+            Context.DEFERRED_IMPORT_RESOURCE_PATH_KEY,
+            interpreter
+              .getContext()
+              .getPyishObjectMapper()
+              .getAsPyishString(importFile.get())
+          ),
+          interpreter,
+          false
+        );
+      suffix =
+        EagerTagDecorator.buildSetTagForDeferredInChildContext(
+          ImmutableMap.of(
+            Context.DEFERRED_IMPORT_RESOURCE_PATH_KEY,
+            interpreter
+              .getContext()
+              .getPyishObjectMapper()
+              .getAsPyishString(currentDeferredImportResource)
+          ),
+          interpreter,
+          false
+        );
     }
 
     String result;
@@ -161,39 +153,4 @@ public String reconstructImage() {
     }
     return prefix + result + suffix;
   }
-
-  private String getAliasedImportResourcePrefix(String importResourcePath) {
-    String prefix = "";
-    String importAlias = fullName.split("\\.", 2)[0];
-    Object aliasMap = interpreter.getContext().get(importAlias);
-    if (aliasMap instanceof DeferredValue) {
-      PyMap deferredAliasMap = new PyMap(new HashMap<>());
-      deferredAliasMap.put(Context.IMPORT_RESOURCE_PATH_KEY, importResourcePath);
-      prefix =
-        EagerTagDecorator.buildDoUpdateTag(
-          importAlias,
-          interpreter
-            .getContext()
-            .getPyishObjectMapper()
-            .getAsPyishString(deferredAliasMap),
-          interpreter
-        );
-    } else if (aliasMap instanceof Map) {
-      PyMap deferredAliasMap = new PyMap(new HashMap<>());
-      deferredAliasMap.put(Context.IMPORT_RESOURCE_PATH_KEY, importResourcePath);
-      prefix =
-        EagerTagDecorator.buildSetTagForDeferredInChildContext(
-          ImmutableMap.of(
-            importAlias,
-            interpreter
-              .getContext()
-              .getPyishObjectMapper()
-              .getAsPyishString(deferredAliasMap)
-          ),
-          interpreter,
-          false
-        );
-    }
-    return prefix;
-  }
 }
diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/MacroTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/MacroTag.java
@@ -14,7 +14,6 @@
 import com.hubspot.jinjava.interpret.TemplateSyntaxException;
 import com.hubspot.jinjava.lib.fn.MacroFunction;
 import com.hubspot.jinjava.tree.TagNode;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -123,34 +122,16 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
     MacroFunction macro;
     String contextImportResourcePath = (String) interpreter
       .getContext()
-      .get(Context.IMPORT_RESOURCE_PATH_KEY, "");
-    String stackImportResourcePath = interpreter
-      .getContext()
-      .getCurrentPathStack()
-      .peek()
-      .orElse("");
+      .get(Context.DEFERRED_IMPORT_RESOURCE_PATH_KEY, "");
+    boolean scopeEntered = false;
     try {
-      if (!contextImportResourcePath.equals(stackImportResourcePath)) {
+      if (StringUtils.isNotEmpty(contextImportResourcePath)) {
+        scopeEntered = true;
         interpreter.enterScope();
         interpreter
           .getContext()
           .put(Context.IMPORT_RESOURCE_PATH_KEY, contextImportResourcePath);
       }
-      if (StringUtils.isNotEmpty(parentName)) {
-        interpreter.enterScope();
-        Object parentAliasMap = interpreter
-          .getContext()
-          .get(parentName, Collections.emptyMap());
-        if (parentAliasMap instanceof DeferredValue) {
-          parentAliasMap = ((DeferredValue) parentAliasMap).getOriginalValue();
-        }
-        String parentImportResourcePath = (String) (
-          (Map<String, Object>) parentAliasMap
-        ).get(Context.IMPORT_RESOURCE_PATH_KEY);
-        interpreter
-          .getContext()
-          .put(Context.IMPORT_RESOURCE_PATH_KEY, parentImportResourcePath);
-      }
       macro =
         new MacroFunction(
           tagNode.getChildren(),
@@ -162,10 +143,7 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
           interpreter.getPosition()
         );
     } finally {
-      if (
-        StringUtils.isNotEmpty(parentName) ||
-        !contextImportResourcePath.equals(stackImportResourcePath)
-      ) {
+      if (scopeEntered) {
         interpreter.leaveScope();
       }
     }
diff --git a/src/test/java/com/hubspot/jinjava/lib/tag/eager/EagerImportTagTest.java b/src/test/java/com/hubspot/jinjava/lib/tag/eager/EagerImportTagTest.java
@@ -470,10 +470,10 @@ public void itCorrectlySetsAliasedPathForSecondPass() {
     );
     assertThat(firstPassResult)
       .isEqualTo(
-        "{% set m = {'import_resource_path': 'import-macro.jinja'} %}{% macro m.print_path_macro(var) %}\n" +
+        "{% set deferred_import_resource_path = 'import-macro.jinja' %}{% macro m.print_path_macro(var) %}\n" +
         "{{ var|print_path }}\n" +
         "{{ var }}\n" +
-        "{% endmacro %}{{ m.print_path_macro(foo) }}"
+        "{% endmacro %}{% set deferred_import_resource_path = null %}{{ m.print_path_macro(foo) }}"
       );
     context.put("foo", "foo");
     assertThat(interpreter.render(firstPassResult).trim())
@@ -489,10 +489,10 @@ public void itCorrectlySetsPathForSecondPass() {
     );
     assertThat(firstPassResult)
       .isEqualTo(
-        "{% set import_resource_path = 'import-macro.jinja' %}{% macro print_path_macro(var) %}\n" +
+        "{% set deferred_import_resource_path = 'import-macro.jinja' %}{% macro print_path_macro(var) %}\n" +
         "{{ var|print_path }}\n" +
         "{{ var }}\n" +
-        "{% endmacro %}{% set import_resource_path = '' %}{{ print_path_macro(foo) }}"
+        "{% endmacro %}{% set deferred_import_resource_path = null %}{{ print_path_macro(foo) }}"
       );
     context.put("foo", "foo");
     assertThat(interpreter.render(firstPassResult).trim())
diff --git a/src/test/resources/eager/handles-deferred-import-vars.expected.jinja b/src/test/resources/eager/handles-deferred-import-vars.expected.jinja
@@ -8,5 +8,5 @@ bar: {{ bar }}
 {% set bar = myname + 19 %}{% set simple = {} %}{% do simple.update({'bar': bar}) %}
 Hello {{ myname }}
 {% do simple.update({'import_resource_path': 'macro-and-set.jinja'}) %}
-simple.foo: {% do simple.update({'import_resource_path': 'macro-and-set.jinja'}) %}{% macro simple.foo() %}Hello {{ myname }}{% endmacro %}{{ simple.foo() }}
+simple.foo: {% set deferred_import_resource_path = 'macro-and-set.jinja' %}{% macro simple.foo() %}Hello {{ myname }}{% endmacro %}{% set deferred_import_resource_path = null %}{{ simple.foo() }}
 simple.bar: {{ simple.bar }}
diff --git a/src/test/resources/eager/puts-deferred-fromed-macro-in-output.expected.jinja b/src/test/resources/eager/puts-deferred-fromed-macro-in-output.expected.jinja
@@ -1,2 +1,2 @@
 {% set myname = deferred + 3 %}
-{% set import_resource_path = 'simple-with-call.jinja' %}{% macro getPath() %}Hello {{ myname }}{% endmacro %}{% set import_resource_path = '' %}{% print getPath() %}
+{% set deferred_import_resource_path = 'simple-with-call.jinja' %}{% macro getPath() %}Hello {{ myname }}{% endmacro %}{% set deferred_import_resource_path = null %}{% print getPath() %}
diff --git a/src/test/resources/eager/puts-deferred-imported-macro-in-output.expected.jinja b/src/test/resources/eager/puts-deferred-imported-macro-in-output.expected.jinja
@@ -1,2 +1,2 @@
 {% set myname = deferred + 3 %}
-{% set simple = {'import_resource_path': 'simple-with-call.jinja'} %}{% macro simple.getPath() %}Hello {{ myname }}{% endmacro %}{% print simple.getPath() %}
+{% set deferred_import_resource_path = 'simple-with-call.jinja' %}{% macro simple.getPath() %}Hello {{ myname }}{% endmacro %}{% set deferred_import_resource_path = null %}{% print simple.getPath() %}

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`{% set myname = deferred + 3 %}`
`2`		`-{% set import_resource_path = 'simple-with-call.jinja' %}{% macro getPath() %}Hello {{ myname }}{% endmacro %}{% set import_resource_path = '' %}{% print getPath() %}`
	`2`	`+{% set deferred_import_resource_path = 'simple-with-call.jinja' %}{% macro getPath() %}Hello {{ myname }}{% endmacro %}{% set deferred_import_resource_path = null %}{% print getPath() %}`
Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`{% set myname = deferred + 3 %}`
`2`		`-{% set simple = {'import_resource_path': 'simple-with-call.jinja'} %}{% macro simple.getPath() %}Hello {{ myname }}{% endmacro %}{% print simple.getPath() %}`
	`2`	`+{% set deferred_import_resource_path = 'simple-with-call.jinja' %}{% macro simple.getPath() %}Hello {{ myname }}{% endmacro %}{% set deferred_import_resource_path = null %}{% print simple.getPath() %}`