From 4a0173101253fdd0025883b6b2ee12cc88418eff Mon Sep 17 00:00:00 2001
From: Jack Smith <jasmith@hubspot.com>
Date: Thu, 5 Oct 2023 11:15:50 -0400
Subject: [PATCH] Support aliasing macro function names in from tag For
 example: ``` {% from 'file.jinja' import my_macro as m %} ```

---
 .../com/hubspot/jinjava/lib/tag/FromTag.java  |  6 ++-
 .../jinjava/lib/tag/eager/EagerFromTag.java   | 38 ++++++-------------
 .../hubspot/jinjava/lib/tag/FromTagTest.java  |  8 ++++
 .../tags/macrotag/from-alias-macro.jinja      |  4 ++
 4 files changed, 28 insertions(+), 28 deletions(-)
 create mode 100644 src/test/resources/tags/macrotag/from-alias-macro.jinja

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/FromTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/FromTag.java
index 562a02de7..feac95fac 100644
--- a/src/main/java/com/hubspot/jinjava/lib/tag/FromTag.java
+++ b/src/main/java/com/hubspot/jinjava/lib/tag/FromTag.java
@@ -167,7 +167,11 @@ public static boolean integrateChild(
       Object val = child.getContext().getGlobalMacro(importMapping.getKey());
 
       if (val != null) {
-        interpreter.getContext().addGlobalMacro((MacroFunction) val);
+        MacroFunction toImport = (MacroFunction) val;
+        if (!importMapping.getKey().equals(importMapping.getValue())) {
+          toImport = new MacroFunction(toImport, importMapping.getValue());
+        }
+        interpreter.getContext().addGlobalMacro(toImport);
       } else {
         val = child.getContext().get(importMapping.getKey());
 
diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerFromTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerFromTag.java
index 9d39f4052..284dc22f4 100644
--- a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerFromTag.java
+++ b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerFromTag.java
@@ -13,11 +13,10 @@
 import com.hubspot.jinjava.tree.parse.TagToken;
 import com.hubspot.jinjava.util.EagerReconstructionUtils;
 import java.io.IOException;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Optional;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 @Beta
@@ -100,11 +99,10 @@ public String getEagerTagImage(TagToken tagToken, JinjavaInterpreter interpreter
         }
 
         FromTag.integrateChild(imports, child, interpreter);
-        Map<String, String> newToOldImportNames = renameMacros(imports, interpreter)
-          .entrySet()
-          .stream()
-          .filter(e -> !e.getKey().equals(e.getValue()))
-          .collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
+        Map<String, String> newToOldImportNames = getNewToOldWithoutMacros(
+          imports,
+          interpreter
+        );
         if (child.getContext().getDeferredTokens().isEmpty() || output == null) {
           return "";
         } else if (newToOldImportNames.size() > 0) {
@@ -132,33 +130,19 @@ public String getEagerTagImage(TagToken tagToken, JinjavaInterpreter interpreter
     }
   }
 
-  private static Map<String, String> renameMacros(
+  private static Map<String, String> getNewToOldWithoutMacros(
     Map<String, String> oldToNewImportNames,
     JinjavaInterpreter interpreter
   ) {
-    Set<String> toRemove = new HashSet<>();
-    Map<String, MacroFunction> macroFunctions = oldToNewImportNames
+    return oldToNewImportNames
       .entrySet()
       .stream()
+      .filter(e -> !e.getKey().equals(e.getValue()))
       .filter(
         e ->
-          !e.getKey().equals(e.getValue()) &&
-          !interpreter.getContext().containsKey(e.getKey()) &&
-          interpreter.getContext().isGlobalMacro(e.getKey())
+          interpreter.getContext().containsKey(e.getValue()) ||
+          !interpreter.getContext().isGlobalMacro(e.getValue())
       )
-      .peek(entry -> toRemove.add(entry.getKey()))
-      .collect(
-        Collectors.toMap(
-          Map.Entry::getValue,
-          e -> interpreter.getContext().getGlobalMacro(e.getKey())
-        )
-      );
-
-    macroFunctions.forEach(
-      (key, value) ->
-        interpreter.getContext().addGlobalMacro(new MacroFunction(value, key))
-    );
-    toRemove.forEach(oldToNewImportNames::remove);
-    return oldToNewImportNames;
+      .collect(Collectors.toMap(Entry::getValue, Entry::getKey)); // flip order
   }
 }
diff --git a/src/test/java/com/hubspot/jinjava/lib/tag/FromTagTest.java b/src/test/java/com/hubspot/jinjava/lib/tag/FromTagTest.java
index e8a386381..6a601c798 100644
--- a/src/test/java/com/hubspot/jinjava/lib/tag/FromTagTest.java
+++ b/src/test/java/com/hubspot/jinjava/lib/tag/FromTagTest.java
@@ -58,6 +58,14 @@ public void importedContextExposesVars() {
       .contains("wrap-padding: padding-left:42px;padding-right:42px");
   }
 
+  @Test
+  public void itImportsAliasedMacroName() {
+    assertThat(fixture("from-alias-macro"))
+      .contains("wrap-spacer:")
+      .contains("<td height=\"42\">")
+      .contains("wrap-padding: padding-left:42px;padding-right:42px");
+  }
+
   @Test
   public void importedCycleDected() {
     fixture("from-recursion");
diff --git a/src/test/resources/tags/macrotag/from-alias-macro.jinja b/src/test/resources/tags/macrotag/from-alias-macro.jinja
new file mode 100644
index 000000000..fdfb0ed21
--- /dev/null
+++ b/src/test/resources/tags/macrotag/from-alias-macro.jinja
@@ -0,0 +1,4 @@
+{% from "pegasus-importable.jinja" import wrap_padding as wp, spacer as sp %}
+
+wrap-padding: {{ wp }}
+wrap-spacer: {{ sp() }}
\ No newline at end of file