PythonTreeMaker { public ReturnStatement returnStatement(StatementWithSeparator statementWithSeparator) { AstNode astNode = statementWithSeparator.statement(); Separators separators = statementWithSeparator.separator(); AstNode testListNode = astNode.getFirstChild(PythonGrammar.TESTLIST_STAR_EXPR); List<Expression> expressionTrees = Collections.emptyList(); List<Token> commas = Collections.emptyList(); if (testListNode != null) { expressionTrees = expressionsFromTestListStarExpr(testListNode); commas = punctuators(testListNode, PythonPunctuator.COMMA); } return new ReturnStatementImpl(toPyToken(astNode.getTokens()).get(0), expressionTrees, commas, separators); }  FileInput fileInput(AstNode astNode); void setParents(Tree root); AnnotatedAssignment annotatedAssignment(StatementWithSeparator statementWithSeparator); PrintStatement printStatement(StatementWithSeparator statementWithSeparator); ExecStatement execStatement(StatementWithSeparator statementWithSeparator); AssertStatement assertStatement(StatementWithSeparator statementWithSeparator); PassStatement passStatement(StatementWithSeparator statementWithSeparator); DelStatement delStatement(StatementWithSeparator statementWithSeparator); ReturnStatement returnStatement(StatementWithSeparator statementWithSeparator); YieldStatement yieldStatement(StatementWithSeparator statementWithSeparator); YieldExpression yieldExpression(AstNode astNode); RaiseStatement raiseStatement(StatementWithSeparator statementWithSeparator); BreakStatement breakStatement(StatementWithSeparator statementWithSeparator); ContinueStatement continueStatement(StatementWithSeparator statementWithSeparator); ImportStatement importStatement(StatementWithSeparator statementWithSeparator); GlobalStatement globalStatement(StatementWithSeparator statementWithSeparator); NonlocalStatement nonlocalStatement(StatementWithSeparator statementWithSeparator); IfStatement ifStatement(AstNode astNode); FunctionDef funcDefStatement(AstNode astNode); ClassDef classDefStatement(AstNode astNode); ForStatement forStatement(AstNode astNode); WhileStatementImpl whileStatement(AstNode astNode); ExpressionStatement expressionStatement(StatementWithSeparator statementWithSeparator); AssignmentStatement assignment(StatementWithSeparator statementWithSeparator); CompoundAssignmentStatement compoundAssignment(StatementWithSeparator statementWithSeparator); TryStatement tryStatement(AstNode astNode); WithStatement withStatement(AstNode astNode); Expression expression(AstNode astNode); Argument argument(AstNode astNode); LambdaExpression lambdaExpression(AstNode astNode);  }
@Test public void returnStatement() { setRootRule(PythonGrammar.RETURN_STMT); AstNode astNode = p.parse("return foo"); StatementWithSeparator statementWithSeparator = new StatementWithSeparator(astNode, null); ReturnStatement returnStatement = treeMaker.returnStatement(statementWithSeparator); assertThat(returnStatement).isNotNull(); assertThat(returnStatement.returnKeyword().value()).isEqualTo("return"); assertThat(returnStatement.expressions()).hasSize(1); assertThat(returnStatement.commas()).isEmpty(); assertThat(returnStatement.children()).hasSize(2); astNode = p.parse("return foo, bar"); statementWithSeparator = new StatementWithSeparator(astNode, null); returnStatement = treeMaker.returnStatement(statementWithSeparator); assertThat(returnStatement).isNotNull(); assertThat(returnStatement.returnKeyword().value()).isEqualTo("return"); assertThat(returnStatement.expressions()).hasSize(2); assertThat(returnStatement.commas()).hasSize(1); assertThat(returnStatement.children()).hasSize(4); assertThat(returnStatement.children().get(2).is(Kind.TOKEN)).isTrue(); astNode = p.parse("return"); statementWithSeparator = new StatementWithSeparator(astNode, null); returnStatement = treeMaker.returnStatement(statementWithSeparator); assertThat(returnStatement).isNotNull(); assertThat(returnStatement.returnKeyword().value()).isEqualTo("return"); assertThat(returnStatement.expressions()).isEmpty(); assertThat(returnStatement.commas()).isEmpty(); assertThat(returnStatement.children()).hasSize(1); astNode = p.parse("return []"); statementWithSeparator = new StatementWithSeparator(astNode, null); returnStatement = treeMaker.returnStatement(statementWithSeparator); ListLiteral listLiteral = (ListLiteral) returnStatement.expressions().get(0); assertThat(listLiteral.leftBracket()).isSameAs(listLiteral.children().get(0)); assertThat(listLiteral.rightBracket()).isSameAs(listLiteral.children().get(2)); ExpressionList emptyExpressionList = listLiteral.elements(); assertThat(emptyExpressionList.children()).isEmpty(); assertThat(emptyExpressionList.firstToken()).isNull(); assertThat(emptyExpressionList.lastToken()).isNull(); astNode = p.parse("return foo(), *bar"); statementWithSeparator = new StatementWithSeparator(astNode, null); returnStatement = treeMaker.returnStatement(statementWithSeparator); assertThat(returnStatement.expressions()).extracting(Tree::getKind).containsExactly(Kind.CALL_EXPR, Kind.UNPACKING_EXPR); }