TreeUtils { public static List<FunctionDef> topLevelFunctionDefs(ClassDef classDef) { CollectFunctionDefsVisitor visitor = new CollectFunctionDefsVisitor(); classDef.body().accept(visitor); return visitor.functionDefs; } 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 topLevelFunctionDefs() { FileInput fileInput = PythonTestUtils.parse( "class A:", " x = True", " def foo(self): pass", " if x:", " def bar(self, x): return 1", " else:", " def baz(self, x, y): return x + y" ); ClassDef classDef = PythonTestUtils.getLastDescendant(fileInput, t -> t.is(Kind.CLASSDEF)); List<FunctionDef> functionDefs = PythonTestUtils.getAllDescendant(fileInput, t -> t.is(Kind.FUNCDEF)); assertThat(TreeUtils.topLevelFunctionDefs(classDef)).containsAll(functionDefs); fileInput = PythonTestUtils.parse( "class A:", " x = True", " def foo(self):", " def foo2(x, y): return x + y", " return foo2(1, 1)", " class B:", " def bar(self): pass" ); classDef = PythonTestUtils.getFirstChild(fileInput, t -> t.is(Kind.CLASSDEF)); FunctionDef fooDef = PythonTestUtils.getLastDescendant(fileInput, t -> t.is(Kind.FUNCDEF) && ((FunctionDef) t).name().name().equals("foo")); assertThat(TreeUtils.topLevelFunctionDefs(classDef)).isEqualTo(Collections.singletonList(fooDef)); }