def test_statement(self) -> None: # Test that we can insert various types of statements into a # statement list. module = parse_template_module( "{statement1}\n{statement2}\n{statement3}\n", statement1=cst.If( test=cst.Name("foo"), body=cst.SimpleStatementSuite((cst.Pass(),),), ), statement2=cst.SimpleStatementLine((cst.Expr(cst.Call(cst.Name("bar"))),),), statement3=cst.Pass(), ) self.assertEqual( module.code, "if foo: pass\nbar()\npass\n", )
def test_module_config_for_parsing(self) -> None: module = parse_module("pass\r") statement = parse_statement("if True:\r pass", config=module.config_for_parsing) self.assertEqual( statement, cst.If( test=cst.Name(value="True"), body=cst.IndentedBlock( body=[cst.SimpleStatementLine(body=[cst.Pass()])], header=cst.TrailingWhitespace(newline=cst.Newline( # This would be "\r" if we didn't pass the module config forward. value=None)), ), ), )
def _get_assert_replacement(self, node: cst.Assert): message = node.msg or str(cst.Module(body=[node]).code) return cst.If( test=cst.UnaryOperation( operator=cst.Not(), expression=node.test, # Todo: parenthesize? ), body=cst.IndentedBlock(body=[ cst.SimpleStatementLine(body=[ cst.Raise(exc=cst.Call( func=cst.Name(value="AssertionError", ), args=[ cst.Arg(value=cst.SimpleString(value=repr(message), ), ), ], ), ), ]), ], ), )
class FooterBehaviorTest(UnitTest): @data_provider({ # Literally the most basic example "simple_module": { "code": "\n", "expected_module": cst.Module(body=()) }, # A module with a header comment "header_only_module": { "code": "# This is a header comment\n", "expected_module": cst.Module( header=[ cst.EmptyLine(comment=cst.Comment( value="# This is a header comment")) ], body=[], ), }, # A module with a header and footer "simple_header_footer_module": { "code": "# This is a header comment\npass\n# This is a footer comment\n", "expected_module": cst.Module( header=[ cst.EmptyLine(comment=cst.Comment( value="# This is a header comment")) ], body=[cst.SimpleStatementLine([cst.Pass()])], footer=[ cst.EmptyLine(comment=cst.Comment( value="# This is a footer comment")) ], ), }, # A module which should have a footer comment taken from the # if statement's indented block. "simple_reparented_footer_module": { "code": "# This is a header comment\nif True:\n pass\n# This is a footer comment\n", "expected_module": cst.Module( header=[ cst.EmptyLine(comment=cst.Comment( value="# This is a header comment")) ], body=[ cst.If( test=cst.Name(value="True"), body=cst.IndentedBlock( header=cst.TrailingWhitespace(), body=[ cst.SimpleStatementLine( body=[cst.Pass()], trailing_whitespace=cst.TrailingWhitespace( ), ) ], ), ) ], footer=[ cst.EmptyLine(comment=cst.Comment( value="# This is a footer comment")) ], ), }, # Verifying that we properly parse and spread out footer comments to the # relative indents they go with. "complex_reparented_footer_module": { "code": ("# This is a header comment\nif True:\n if True:\n pass" + "\n # This is an inner indented block comment\n # This " + "is an outer indented block comment\n# This is a footer comment\n" ), "expected_module": cst.Module( body=[ cst.If( test=cst.Name(value="True"), body=cst.IndentedBlock( body=[ cst.If( test=cst.Name(value="True"), body=cst.IndentedBlock( body=[ cst.SimpleStatementLine( body=[cst.Pass()]) ], footer=[ cst.EmptyLine(comment=cst.Comment( value= "# This is an inner indented block comment" )) ], ), ) ], footer=[ cst.EmptyLine(comment=cst.Comment( value= "# This is an outer indented block comment" )) ], ), ) ], header=[ cst.EmptyLine(comment=cst.Comment( value="# This is a header comment")) ], footer=[ cst.EmptyLine(comment=cst.Comment( value="# This is a footer comment")) ], ), }, # Verify that comments belonging to statements are still owned even # after an indented block. "statement_comment_reparent": { "code": "if foo:\n return\n# comment\nx = 7\n", "expected_module": cst.Module(body=[ cst.If( test=cst.Name(value="foo"), body=cst.IndentedBlock(body=[ cst.SimpleStatementLine(body=[ cst.Return( whitespace_after_return=cst.SimpleWhitespace( value="")) ]) ]), ), cst.SimpleStatementLine( body=[ cst.Assign( targets=[ cst.AssignTarget(target=cst.Name(value="x")) ], value=cst.Integer(value="7"), ) ], leading_lines=[ cst.EmptyLine(comment=cst.Comment(value="# comment")) ], ), ]), }, # Verify that even if there are completely empty lines, we give all lines # up to and including the last line that's indented correctly. That way # comments that line up with indented block's indentation level aren't # parented to the next line just because there's a blank line or two # between them. "statement_comment_with_empty_lines": { "code": ("def foo():\n if True:\n pass\n\n # Empty " + "line before me\n\n else:\n pass\n"), "expected_module": cst.Module(body=[ cst.FunctionDef( name=cst.Name(value="foo"), params=cst.Parameters(), body=cst.IndentedBlock(body=[ cst.If( test=cst.Name(value="True"), body=cst.IndentedBlock( body=[ cst.SimpleStatementLine(body=[cst.Pass()]) ], footer=[ cst.EmptyLine(indent=False), cst.EmptyLine(comment=cst.Comment( value="# Empty line before me")), ], ), orelse=cst.Else( body=cst.IndentedBlock(body=[ cst.SimpleStatementLine(body=[cst.Pass()]) ]), leading_lines=[cst.EmptyLine(indent=False)], ), ) ]), ) ]), }, }) def test_parsers(self, code: str, expected_module: cst.CSTNode) -> None: parsed_module = parse_module(dedent(code)) self.assertTrue( deep_equals(parsed_module, expected_module), msg= f"\n{parsed_module!r}\nis not deeply equal to \n{expected_module!r}", )
class IfTest(CSTNodeTest): @data_provider(( # Simple if without elif or else { "node": cst.If(cst.Name("conditional"), cst.SimpleStatementSuite((cst.Pass(), ))), "code": "if conditional: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (1, 20)), }, # else clause { "node": cst.If( cst.Name("conditional"), cst.SimpleStatementSuite((cst.Pass(), )), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(), ))), ), "code": "if conditional: pass\nelse: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (2, 10)), }, # elif clause { "node": cst.If( cst.Name("conditional"), cst.SimpleStatementSuite((cst.Pass(), )), orelse=cst.If( cst.Name("other_conditional"), cst.SimpleStatementSuite((cst.Pass(), )), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(), ))), ), ), "code": "if conditional: pass\nelif other_conditional: pass\nelse: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (3, 10)), }, # indentation { "node": DummyIndentedBlock( " ", cst.If( cst.Name("conditional"), cst.SimpleStatementSuite((cst.Pass(), )), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(), ))), ), ), "code": " if conditional: pass\n else: pass\n", "parser": None, "expected_position": CodeRange((1, 4), (2, 14)), }, # with an indented body { "node": DummyIndentedBlock( " ", cst.If( cst.Name("conditional"), cst.IndentedBlock((cst.SimpleStatementLine( (cst.Pass(), )), )), ), ), "code": " if conditional:\n pass\n", "parser": None, "expected_position": CodeRange((1, 4), (2, 12)), }, # leading_lines { "node": cst.If( cst.Name("conditional"), cst.SimpleStatementSuite((cst.Pass(), )), leading_lines=(cst.EmptyLine( comment=cst.Comment("# leading comment")), ), ), "code": "# leading comment\nif conditional: pass\n", "parser": parse_statement, "expected_position": CodeRange((2, 0), (2, 20)), }, # whitespace before/after test and else { "node": cst.If( cst.Name("conditional"), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_before_test=cst.SimpleWhitespace(" "), whitespace_after_test=cst.SimpleWhitespace(" "), orelse=cst.Else( cst.SimpleStatementSuite((cst.Pass(), )), whitespace_before_colon=cst.SimpleWhitespace(" "), ), ), "code": "if conditional : pass\nelse : pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (2, 11)), }, # empty lines between if/elif/else clauses, not captured by the suite. { "node": cst.If( cst.Name("test_a"), cst.SimpleStatementSuite((cst.Pass(), )), orelse=cst.If( cst.Name("test_b"), cst.SimpleStatementSuite((cst.Pass(), )), leading_lines=(cst.EmptyLine(), ), orelse=cst.Else( cst.SimpleStatementSuite((cst.Pass(), )), leading_lines=(cst.EmptyLine(), ), ), ), ), "code": "if test_a: pass\n\nelif test_b: pass\n\nelse: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (5, 10)), }, )) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs)
def leave_Module(self, original_node: libcst.Module, updated_node: libcst.Module) -> libcst.Module: # Don't try to modify if we have nothing to do if (not self.module_imports and not self.module_mapping and not self.module_aliases and not self.alias_mapping): return updated_node # First, find the insertion point for imports ( statements_before_imports, statements_until_add_imports, statements_after_imports, ) = self._split_module(original_node, updated_node) # Make sure there's at least one empty line before the first non-import statements_after_imports = self._insert_empty_line( statements_after_imports) # Mapping of modules we're adding to the object with and without alias they should import module_and_alias_mapping = defaultdict(list) for module, aliases in self.alias_mapping.items(): module_and_alias_mapping[module].extend(aliases) for module, imports in self.module_mapping.items(): module_and_alias_mapping[module].extend([(object, None) for object in imports]) module_and_alias_mapping = { module: sorted(aliases) for module, aliases in module_and_alias_mapping.items() } import_cycle_safe_module_names = [ 'mypy_extensions', 'typing', 'typing_extensions', ] type_checking_cond_import = parse_statement( f"from typing import TYPE_CHECKING", config=updated_node.config_for_parsing, ) type_checking_cond_statement = libcst.If( test=libcst.Name("TYPE_CHECKING"), body=libcst.IndentedBlock(body=[ parse_statement( f"from {module} import " + ", ".join([ obj if alias is None else f"{obj} as {alias}" for (obj, alias) in aliases ]), config=updated_node.config_for_parsing, ) for module, aliases in module_and_alias_mapping.items() if module != "__future__" and module not in import_cycle_safe_module_names ], ), ) if not type_checking_cond_statement.body.body: type_checking_cond_statement = libcst.EmptyLine() type_checking_cond_import = libcst.EmptyLine() # import ptvsd; ptvsd.set_trace() # Now, add all of the imports we need! return updated_node.with_changes(body=( *statements_before_imports, *[ parse_statement( f"from {module} import " + ", ".join([ obj if alias is None else f"{obj} as {alias}" for (obj, alias) in aliases ]), config=updated_node.config_for_parsing, ) for module, aliases in module_and_alias_mapping.items() if module == "__future__" ], *statements_until_add_imports, *[ parse_statement(f"import {module}", config=updated_node.config_for_parsing) for module in sorted(self.module_imports) ], *[ parse_statement( f"import {module} as {asname}", config=updated_node.config_for_parsing, ) for (module, asname) in self.module_aliases.items() ], # TODO: 可以进一步用 `from __future__ import annotations` 解决forward ref, 这里加也可以,用其他工具也可以 type_checking_cond_import, type_checking_cond_statement, *[ parse_statement( f"from {module} import " + ", ".join([ obj if alias is None else f"{obj} as {alias}" for (obj, alias) in aliases ]), config=updated_node.config_for_parsing, ) for module, aliases in module_and_alias_mapping.items() if module != "__future__" and module in import_cycle_safe_module_names and not module.startswith("monkeytype") ], *statements_after_imports, ))
class NamedExprTest(CSTNodeTest): @data_provider(( # Simple named expression { "node": cst.NamedExpr(cst.Name("x"), cst.Float("5.5")), "code": "x := 5.5", "parser": None, # Walrus operator is illegal as top-level statement "expected_position": None, }, # Parenthesized named expression { "node": cst.NamedExpr( lpar=(cst.LeftParen(), ), target=cst.Name("foo"), value=cst.Integer("5"), rpar=(cst.RightParen(), ), ), "code": "(foo := 5)", "parser": _parse_expression_force_38, "expected_position": CodeRange((1, 1), (1, 9)), }, # Make sure that spacing works { "node": cst.NamedExpr( lpar=(cst.LeftParen( whitespace_after=cst.SimpleWhitespace(" ")), ), target=cst.Name("foo"), whitespace_before_walrus=cst.SimpleWhitespace(" "), whitespace_after_walrus=cst.SimpleWhitespace(" "), value=cst.Name("bar"), rpar=(cst.RightParen( whitespace_before=cst.SimpleWhitespace(" ")), ), ), "code": "( foo := bar )", "parser": _parse_expression_force_38, "expected_position": CodeRange((1, 2), (1, 14)), }, # Make sure we can use these where allowed in if/while statements { "node": cst.While( test=cst.NamedExpr( target=cst.Name(value="x"), value=cst.Call(func=cst.Name(value="some_input")), ), body=cst.SimpleStatementSuite(body=[cst.Pass()]), ), "code": "while x := some_input(): pass\n", "parser": _parse_statement_force_38, "expected_position": None, }, { "node": cst.If( test=cst.NamedExpr( target=cst.Name(value="x"), value=cst.Call(func=cst.Name(value="some_input")), ), body=cst.SimpleStatementSuite(body=[cst.Pass()]), ), "code": "if x := some_input(): pass\n", "parser": _parse_statement_force_38, "expected_position": None, }, { "node": cst.If( test=cst.NamedExpr( target=cst.Name(value="x"), value=cst.Integer(value="1"), whitespace_before_walrus=cst.SimpleWhitespace(""), whitespace_after_walrus=cst.SimpleWhitespace(""), ), body=cst.SimpleStatementSuite(body=[cst.Pass()]), ), "code": "if x:=1: pass\n", "parser": _parse_statement_force_38, "expected_position": None, }, # Function args { "node": cst.Call( func=cst.Name(value="f"), args=[ cst.Arg(value=cst.NamedExpr( target=cst.Name(value="y"), value=cst.Integer(value="1"), whitespace_before_walrus=cst.SimpleWhitespace(""), whitespace_after_walrus=cst.SimpleWhitespace(""), )), ], ), "code": "f(y:=1)", "parser": _parse_expression_force_38, "expected_position": None, }, # Whitespace handling on args is fragile { "node": cst.Call( func=cst.Name(value="f"), args=[ cst.Arg( value=cst.Name(value="x"), comma=cst.Comma( whitespace_after=cst.SimpleWhitespace(" ")), ), cst.Arg( value=cst.NamedExpr( target=cst.Name(value="y"), value=cst.Integer(value="1"), whitespace_before_walrus=cst.SimpleWhitespace( " "), whitespace_after_walrus=cst.SimpleWhitespace( " "), ), whitespace_after_arg=cst.SimpleWhitespace(" "), ), ], ), "code": "f(x, y := 1 )", "parser": _parse_expression_force_38, "expected_position": None, }, { "node": cst.Call( func=cst.Name(value="f"), args=[ cst.Arg( value=cst.NamedExpr( target=cst.Name(value="y"), value=cst.Integer(value="1"), whitespace_before_walrus=cst.SimpleWhitespace( " "), whitespace_after_walrus=cst.SimpleWhitespace( " "), ), whitespace_after_arg=cst.SimpleWhitespace(" "), ), ], whitespace_before_args=cst.SimpleWhitespace(" "), ), "code": "f( y := 1 )", "parser": _parse_expression_force_38, "expected_position": None, }, )) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs) @data_provider(( { "get_node": (lambda: cst.NamedExpr( cst.Name("foo"), cst.Name("bar"), lpar=(cst.LeftParen(), ))), "expected_re": "left paren without right paren", }, { "get_node": (lambda: cst.NamedExpr( cst.Name("foo"), cst.Name("bar"), rpar=(cst.RightParen(), ))), "expected_re": "right paren without left paren", }, )) def test_invalid(self, **kwargs: Any) -> None: self.assert_invalid(**kwargs)
class IndentedBlockTest(CSTNodeTest): @data_provider(( # Standard render ( cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(), )), )), "\n pass\n", None, ), # Render with empty (cst.IndentedBlock(()), "\n pass\n", None), # Render with empty subnodes (cst.IndentedBlock((cst.SimpleStatementLine( ()), )), "\n pass\n", None), # Test render with custom indent ( cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(), )), ), indent="\t"), "\n\tpass\n", None, ), # Test comments ( cst.IndentedBlock( (cst.SimpleStatementLine((cst.Pass(), )), ), header=cst.TrailingWhitespace( whitespace=cst.SimpleWhitespace(" "), comment=cst.Comment("# header comment"), ), ), " # header comment\n pass\n", None, ), ( cst.IndentedBlock( (cst.SimpleStatementLine((cst.Pass(), )), ), footer=(cst.EmptyLine( comment=cst.Comment("# footer comment")), ), ), "\n pass\n # footer comment\n", None, ), ( cst.IndentedBlock( (cst.SimpleStatementLine((cst.Pass(), )), ), footer=(cst.EmptyLine( whitespace=cst.SimpleWhitespace(" "), comment=cst.Comment("# footer comment"), ), ), ), "\n pass\n # footer comment\n", None, ), ( cst.IndentedBlock(( cst.SimpleStatementLine((cst.Continue(), )), cst.SimpleStatementLine((cst.Pass(), )), )), "\n continue\n pass\n", None, ), # Basic parsing test ( cst.If( cst.Name("conditional"), cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(), )), )), ), "if conditional:\n pass\n", parse_statement, ), # Multi-level parsing test ( cst.If( cst.Name("conditional"), cst.IndentedBlock(( cst.SimpleStatementLine((cst.Pass(), )), cst.If( cst.Name("other_conditional"), cst.IndentedBlock((cst.SimpleStatementLine( (cst.Pass(), )), )), ), )), ), "if conditional:\n pass\n if other_conditional:\n pass\n", parse_statement, ), # Inconsistent indentation parsing test ( cst.If( cst.Name("conditional"), cst.IndentedBlock(( cst.SimpleStatementLine((cst.Pass(), )), cst.If( cst.Name("other_conditional"), cst.IndentedBlock( (cst.SimpleStatementLine((cst.Pass(), )), ), indent=" ", ), ), )), ), "if conditional:\n pass\n if other_conditional:\n pass\n", parse_statement, ), )) def test_valid( self, node: cst.CSTNode, code: str, parser: Optional[Callable[[str], cst.CSTNode]], ) -> None: self.validate_node(node, code, parser) @data_provider(( ( lambda: cst.IndentedBlock( (cst.SimpleStatementLine((cst.Pass(), )), ), indent=""), "non-zero width indent", ), ( lambda: cst.IndentedBlock( (cst.SimpleStatementLine((cst.Pass(), )), ), indent="this isn't valid whitespace!", ), "only whitespace", ), )) def test_invalid(self, get_node: Callable[[], cst.CSTNode], expected_re: str) -> None: self.assert_invalid(get_node, expected_re)