Add functional visitor methods

javaparser · Oct 17, 2017 · b1bf722 · b1bf722
1 parent 23669ee
commit b1bf722
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 12 deletions.
diff --git a/javaparser-core/src/main/java/com/github/javaparser/ast/Node.java b/javaparser-core/src/main/java/com/github/javaparser/ast/Node.java
@@ -31,22 +31,22 @@
 import com.github.javaparser.ast.observer.AstObserver;
 import com.github.javaparser.ast.observer.ObservableProperty;
 import com.github.javaparser.ast.observer.PropagatingAstObserver;
-import com.github.javaparser.ast.visitor.CloneVisitor;
-import com.github.javaparser.ast.visitor.EqualsVisitor;
-import com.github.javaparser.ast.visitor.HashCodeVisitor;
-import com.github.javaparser.ast.visitor.Visitable;
+import com.github.javaparser.ast.visitor.*;
 import com.github.javaparser.metamodel.InternalProperty;
 import com.github.javaparser.metamodel.JavaParserMetaModel;
 import com.github.javaparser.metamodel.NodeMetaModel;
 import com.github.javaparser.metamodel.PropertyMetaModel;
 import com.github.javaparser.printer.PrettyPrinter;
 import com.github.javaparser.printer.PrettyPrinterConfiguration;
 import com.github.javaparser.resolution.SymbolResolver;
+
 import javax.annotation.Generated;
 import java.util.*;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
+
 import static com.github.javaparser.ast.Node.Parsedness.PARSED;
 import static java.util.Collections.unmodifiableList;
-import com.github.javaparser.ast.Node;
 
 /**
  * Base class for all nodes of the abstract syntax tree.
@@ -409,6 +409,7 @@ public void tryAddImportToParentCompilationUnit(Class<?> clazz) {
      * Recursively finds all nodes of a certain type.
      *
      * @param clazz the type of node to find.
+     * @deprecated use find(Class)
      */
     public <N extends Node> List<N> getChildNodesByType(Class<N> clazz) {
         List<N> nodes = new ArrayList<>();
@@ -422,7 +423,7 @@ public <N extends Node> List<N> getChildNodesByType(Class<N> clazz) {
     }
 
     /**
-     * @deprecated use getChildNodesByType
+     * @deprecated use find(Class)
      */
     @Deprecated
     public <N extends Node> List<N> getNodesByType(Class<N> clazz) {
@@ -677,4 +678,50 @@ protected SymbolResolver getSymbolResolver() {
     // We need to expose it because we will need to use it to inject the SymbolSolver
     public static final DataKey<SymbolResolver> SYMBOL_RESOLVER_KEY = new DataKey<SymbolResolver>() {
     };
+
+    /**
+     * Walks the AST, calling the consumer for every node.
+     */
+    public void walk(Consumer<Node> consumer) {
+        new TreeVisitor() {
+            @Override
+            public void process(Node node) {
+                consumer.accept(node);
+            }
+        }.visitPreOrder(this);
+    }
+
+    /**
+     * Walks the AST, calling the consumer for every node of type "nodeType".
+     */
+    public <T extends Node> void walk(Class<T> nodeType, Consumer<T> consumer) {
+        new TreeVisitor() {
+            @Override
+            public void process(Node node) {
+                if (nodeType.isInstance(node)) {
+                    consumer.accept(nodeType.cast(node));
+                }
+            }
+        }.visitPreOrder(this);
+    }
+
+    /**
+     * Walks the AST, returning the all nodes of type "nodeType".
+     */
+    public <T extends Node> List<T> find(Class<T> nodeType) {
+        final List<T> found = new ArrayList<>();
+        walk(nodeType, found::add);
+        return found;
+    }
+
+    /**
+     * Walks the AST, returning the all nodes of type "nodeType" that match the predicate.
+     */
+    public <T extends Node> List<T> find(Class<T> nodeType, Predicate<T> predicate) {
+        final List<T> found = new ArrayList<>();
+        walk(nodeType, n -> {
+            if (predicate.test(n)) found.add(n);
+        });
+        return found;
+    }
 }
diff --git a/javaparser-testing/src/test/java/com/github/javaparser/ast/NodeTest.java b/javaparser-testing/src/test/java/com/github/javaparser/ast/NodeTest.java
@@ -21,7 +21,6 @@
 
 package com.github.javaparser.ast;
 
-import com.github.javaparser.JavaParser;
 import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
 import com.github.javaparser.ast.body.FieldDeclaration;
 import com.github.javaparser.ast.body.MethodDeclaration;
@@ -30,6 +29,8 @@
 import com.github.javaparser.ast.comments.Comment;
 import com.github.javaparser.ast.comments.JavadocComment;
 import com.github.javaparser.ast.comments.LineComment;
+import com.github.javaparser.ast.expr.Expression;
+import com.github.javaparser.ast.expr.IntegerLiteralExpr;
 import com.github.javaparser.ast.observer.AstObserver;
 import com.github.javaparser.ast.observer.AstObserverAdapter;
 import com.github.javaparser.ast.observer.ObservableProperty;
@@ -42,6 +43,7 @@
 import java.util.List;
 
 import static com.github.javaparser.JavaParser.parse;
+import static com.github.javaparser.JavaParser.parseExpression;
 import static org.junit.Assert.*;
 
 public class NodeTest {
@@ -302,7 +304,7 @@ public void hasJavaDocCommentNegativeCaseBlockComment() {
         decl.setComment(new BlockComment("foo"));
         assertEquals(false, decl.hasJavaDocComment());
     }
-    
+
     @Test
     public void removeAllOnRequiredProperty() {
         CompilationUnit cu = parse("class X{ void x(){}}");
@@ -314,7 +316,7 @@ public void removeAllOnRequiredProperty() {
 
     @Test
     public void removingTheSecondOfAListOfIdenticalStatementsDoesNotMessUpTheParents() {
-        CompilationUnit unit = JavaParser.parse("public class Example {\n" +
+        CompilationUnit unit = parse("public class Example {\n" +
                 "  public static void example() {\n" +
                 "    boolean swapped;\n" +
                 "    swapped=false;\n" +
@@ -327,17 +329,47 @@ public void removingTheSecondOfAListOfIdenticalStatementsDoesNotMessUpTheParents
         // This will throw an exception if the parents are bad.
         System.out.println(unit.toString());
     }
-    
+
     @Test
     public void findCompilationUnit() {
-        CompilationUnit cu = JavaParser.parse("class X{int x;}");
+        CompilationUnit cu = parse("class X{int x;}");
         VariableDeclarator x = cu.getClassByName("X").get().getMember(0).asFieldDeclaration().getVariables().get(0);
         assertEquals(cu, x.findCompilationUnit().get());
     }
-    
+
     @Test
     public void cantFindCompilationUnit() {
         VariableDeclarator x = new VariableDeclarator();
         assertFalse(x.findCompilationUnit().isPresent());
     }
+
+    @Test
+    public void walk1() {
+        Expression e = parseExpression("1+1");
+        StringBuilder b = new StringBuilder();
+        e.walk(n -> b.append(n.toString()));
+        assertEquals("1 + 111", b.toString());
+    }
+
+    @Test
+    public void walk2() {
+        Expression e = parseExpression("1+1");
+        StringBuilder b = new StringBuilder();
+        e.walk(IntegerLiteralExpr.class, n -> b.append(n.toString()));
+        assertEquals("11", b.toString());
+    }
+
+    @Test
+    public void find1() {
+        Expression e = parseExpression("1+2+3");
+        List<IntegerLiteralExpr> ints = e.find(IntegerLiteralExpr.class, n -> n.asInt() > 1);
+        assertEquals("[2, 3]", ints.toString());
+    }
+
+    @Test
+    public void find2() {
+        Expression e = parseExpression("1+2+3");
+        List<IntegerLiteralExpr> ints = e.find(IntegerLiteralExpr.class);
+        assertEquals("[1, 2, 3]", ints.toString());
+    }
 }