From 8485b480d4175c03a3b5f9b8c4ab6eef932f772e Mon Sep 17 00:00:00 2001
From: Matt Coley <matthewjcoley@gmail.com>
Date: Thu, 11 Jun 2020 13:21:04 -0400
Subject: [PATCH] Revert "Revert "Support expressions in selectattr/rejectattr
 again, more efficiently""

---
 .../hubspot/jinjava/interpret/Context.java    |  7 +++
 .../jinjava/lib/filter/SelectAttrFilter.java  | 24 ++++++--
 .../lib/filter/SelectAttrFilterTest.java      | 61 +++++++++++++++++--
 3 files changed, 81 insertions(+), 11 deletions(-)

diff --git a/src/main/java/com/hubspot/jinjava/interpret/Context.java b/src/main/java/com/hubspot/jinjava/interpret/Context.java
index 0afa5d4fb..bea8caefe 100644
--- a/src/main/java/com/hubspot/jinjava/interpret/Context.java
+++ b/src/main/java/com/hubspot/jinjava/interpret/Context.java
@@ -235,6 +235,13 @@ public void addResolvedExpression(String expression) {
     }
   }
 
+  public void removeResolvedExpression(String expression) {
+    resolvedExpressions.remove(expression);
+    if (getParent() != null) {
+      getParent().removeResolvedExpression(expression);
+    }
+  }
+
   public Set<String> getResolvedExpressions() {
     return ImmutableSet.copyOf(resolvedExpressions);
   }
diff --git a/src/main/java/com/hubspot/jinjava/lib/filter/SelectAttrFilter.java b/src/main/java/com/hubspot/jinjava/lib/filter/SelectAttrFilter.java
index 0a03d35dc..cd6353ee5 100644
--- a/src/main/java/com/hubspot/jinjava/lib/filter/SelectAttrFilter.java
+++ b/src/main/java/com/hubspot/jinjava/lib/filter/SelectAttrFilter.java
@@ -10,11 +10,11 @@
 import com.hubspot.jinjava.lib.exptest.ExpTest;
 import com.hubspot.jinjava.util.ForLoop;
 import com.hubspot.jinjava.util.ObjectIterator;
-import com.hubspot.jinjava.util.Variable;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
 
 @JinjavaDoc(
   value = "Filters a sequence of objects by applying a test to an attribute of an object and only selecting the ones with the test succeeding.",
@@ -110,20 +110,32 @@ protected Object applyFilter(
       }
     }
 
+    String tempValue = generateTempVariable();
+    String expression = generateTempVariable(tempValue, attr);
     ForLoop loop = ObjectIterator.getLoop(var);
     while (loop.hasNext()) {
       Object val = loop.next();
+      interpreter.getContext().put(tempValue, val);
+
+      Object attrVal = interpreter.resolveELExpression(
+        expression,
+        interpreter.getLineNumber()
+      );
 
-      Object attrVal = new Variable(
-        interpreter,
-        String.format("%s.%s", "placeholder", attr)
-      )
-      .resolve(val);
       if (acceptObjects == expTest.evaluate(attrVal, interpreter, expArgs)) {
         result.add(val);
       }
     }
+    interpreter.getContext().removeResolvedExpression(expression);
 
     return result;
   }
+
+  private String generateTempVariable() {
+    return "jj_temp_" + Math.abs(ThreadLocalRandom.current().nextInt());
+  }
+
+  private String generateTempVariable(String tempValue, String expression) {
+    return String.format("%s.%s", tempValue, expression).trim();
+  }
 }
diff --git a/src/test/java/com/hubspot/jinjava/lib/filter/SelectAttrFilterTest.java b/src/test/java/com/hubspot/jinjava/lib/filter/SelectAttrFilterTest.java
index 8e604870e..66d16807b 100644
--- a/src/test/java/com/hubspot/jinjava/lib/filter/SelectAttrFilterTest.java
+++ b/src/test/java/com/hubspot/jinjava/lib/filter/SelectAttrFilterTest.java
@@ -2,9 +2,11 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import com.hubspot.jinjava.Jinjava;
 import java.util.HashMap;
+import java.util.List;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -19,9 +21,27 @@ public void setup() {
       .put(
         "users",
         Lists.newArrayList(
-          new User(0, false, "foo@bar.com", new Option(0, "option0")),
-          new User(1, true, "bar@bar.com", new Option(1, "option1")),
-          new User(2, false, null, new Option(2, "option2"))
+          new User(
+            0,
+            false,
+            "foo@bar.com",
+            new Option(0, "option0"),
+            ImmutableList.of(new Option(0, "option0"))
+          ),
+          new User(
+            1,
+            true,
+            "bar@bar.com",
+            new Option(1, "option1"),
+            ImmutableList.of(new Option(1, "option1"))
+          ),
+          new User(
+            2,
+            false,
+            null,
+            new Option(2, "option2"),
+            ImmutableList.of(new Option(2, "option2"))
+          )
         )
       );
   }
@@ -89,17 +109,44 @@ public void selectAttrWithNestedProperty() {
       .isEqualTo("[2]");
   }
 
+  @Test
+  public void selectAttrWithNestedFilter() {
+    assertThat(
+        jinjava.render(
+          "{{ users|selectattr(\"optionList|map('id')\", 'containing', 1) }}",
+          new HashMap<String, Object>()
+        )
+      )
+      .isEqualTo("[1]");
+
+    assertThat(
+        jinjava.render(
+          "{{ users|selectattr(\"optionList|map('name')\", 'containing', 'option2') }}",
+          new HashMap<String, Object>()
+        )
+      )
+      .isEqualTo("[2]");
+  }
+
   public static class User {
     private long num;
     private boolean isActive;
     private String email;
     private Option option;
-
-    public User(long num, boolean isActive, String email, Option option) {
+    private List<Option> optionList;
+
+    public User(
+      long num,
+      boolean isActive,
+      String email,
+      Option option,
+      List<Option> optionList
+    ) {
       this.num = num;
       this.isActive = isActive;
       this.email = email;
       this.option = option;
+      this.optionList = optionList;
     }
 
     public long getNum() {
@@ -118,6 +165,10 @@ public Option getOption() {
       return option;
     }
 
+    public List<Option> getOptionList() {
+      return optionList;
+    }
+
     @Override
     public String toString() {
       return num + "";