TreeUtils { @CheckForNull public static ClassSymbol getClassSymbolFromDef(@Nullable ClassDef classDef) { if (classDef == null) { return null; } Symbol classNameSymbol = classDef.name().symbol(); if (classNameSymbol == null) { throw new IllegalStateException("A ClassDef should always have a non-null symbol!"); } if (classNameSymbol.kind() == Symbol.Kind.CLASS) { return ((ClassSymbol) classNameSymbol); } return null; } 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 getClassSymbolFromDef() { FileInput fileInput = PythonTestUtils.parse("class A:\n def foo(): pass"); ClassDef classDef = PythonTestUtils.getLastDescendant(fileInput, t -> t.is(Kind.CLASSDEF)); Symbol symbolA = classDef.name().symbol(); assertThat(TreeUtils.getClassSymbolFromDef(classDef)).isEqualTo(symbolA); assertThat(TreeUtils.getClassSymbolFromDef(null)).isNull(); fileInput = PythonTestUtils.parse( "class A:", " pass", "A = 42" ); classDef = PythonTestUtils.getLastDescendant(fileInput, t -> t.is(Kind.CLASSDEF)); assertThat(TreeUtils.getClassSymbolFromDef(classDef)).isNull(); }
@Test(expected = IllegalStateException.class) public void getClassSymbolFromDef_illegalSymbol() { FileInput fileInput = PythonTestUtils.parseWithoutSymbols("class A:\n def foo(): pass"); ClassDef classDef = PythonTestUtils.getLastDescendant(fileInput, t -> t.is(Kind.CLASSDEF)); TreeUtils.getClassSymbolFromDef(classDef); }