TreeUtils { public static List<Token> tokens(Tree tree) { if (tree.is(Kind.TOKEN)) { return Collections.singletonList((Token) tree); } else if (tree.is(Kind.STRING_ELEMENT)) { return Collections.singletonList(tree.firstToken()); } List<Token> tokens = new ArrayList<>(); for (Tree child : tree.children()) { if (child.is(Kind.TOKEN)) { tokens.add(((Token) child)); } else { tokens.addAll(tokens(child)); } } return tokens; } private  TreeUtils(); @CheckForNull static Tree firstAncestor(Tree tree, Predicate<Tree> predicate); @CheckForNull static Tree firstAncestorOfKind(Tree tree, Kind... kinds); static List<Token> tokens(Tree tree); static List<Token> nonWhitespaceTokens(Tree tree); static boolean hasDescendant(Tree tree, Predicate<Tree> predicate); static Stream<Expression> flattenTuples(Expression expression); static Optional<Symbol> getSymbolFromTree(@Nullable Tree tree); @CheckForNull static ClassSymbol getClassSymbolFromDef(@Nullable ClassDef classDef); @CheckForNull static FunctionSymbol getFunctionSymbolFromDef(@Nullable FunctionDef functionDef); static List<Parameter> nonTupleParameters(FunctionDef functionDef); static List<Parameter> positionalParameters(FunctionDef functionDef); static List<FunctionDef> topLevelFunctionDefs(ClassDef classDef); @CheckForNull static RegularArgument argumentByKeyword(String keyword, List<Argument> arguments); @CheckForNull static RegularArgument nthArgumentOrKeyword(int argPosition, String keyword, List<Argument> arguments); static boolean isBooleanLiteral(Tree tree); static String nameFromExpression(Expression expression); @CheckForNull static LocationInFile locationInFile(Tree tree, @Nullable String fileId);  }
@Test public void tokens() { FileInput parsed = parse("if foo:\n pass"); IfStatement ifStmt = (IfStatement) parsed.statements().statements().get(0); List<Token> collect = new ArrayList<>(ifStmt.children().stream().map(t -> t.is(Kind.TOKEN) ? (Token) t : t.firstToken()).collect(Collectors.toList())); collect.add(parsed.lastToken()); assertThat(TreeUtils.tokens(parsed)).containsExactly(collect.toArray(new Token[0])); assertThat(TreeUtils.tokens(parsed.lastToken())).containsExactly(parsed.lastToken()); }