def test_replace_sentinel(self) -> None: def _swap_bools( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.Name("True" if cst.ensure_type(node, cst.Name).value == "False" else "False") # Verify behavior when provided a sentinel replaced = m.replace(cst.RemovalSentinel.REMOVE, m.Name("True") | m.Name("False"), _swap_bools) self.assertEqual(replaced, cst.RemovalSentinel.REMOVE) replaced = m.replace(cst.MaybeSentinel.DEFAULT, m.Name("True") | m.Name("False"), _swap_bools) self.assertEqual(replaced, cst.MaybeSentinel.DEFAULT)
def test_replace_updated_node_changes(self) -> None: def _replace_nested( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.ensure_type(node, cst.Call).with_changes(args=[ cst.Arg( cst.Name(value=cst.ensure_type( cst.ensure_type(extraction["inner"], cst.Call).func, cst.Name, ).value + "_immediate")) ]) original = cst.parse_module( "def foo(val: int) -> int:\n return val\nbar = foo\nbaz = foo\nbiz = foo\nfoo(bar(baz(biz(5))))\n" ) replaced = cst.ensure_type( m.replace( original, m.Call(args=[m.Arg(m.SaveMatchedNode(m.Call(), "inner"))]), _replace_nested, ), cst.Module, ).code self.assertEqual( replaced, "def foo(val: int) -> int:\n return val\nbar = foo\nbaz = foo\nbiz = foo\nfoo(bar_immediate)\n", )
def test_replace_metadata(self) -> None: def _rename_foo( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.ensure_type(node, cst.Name).with_changes(value="replaced") original = cst.parse_module( "foo: int = 37\ndef bar(foo: int) -> int:\n return foo\n\nbiz: int = bar(42)\n" ) wrapper = cst.MetadataWrapper(original) replaced = cst.ensure_type( m.replace( wrapper, m.Name(metadata=m.MatchMetadataIfTrue( meta.QualifiedNameProvider, lambda qualnames: any(n.name == "foo" for n in qualnames), )), _rename_foo, ), cst.Module, ).code self.assertEqual( replaced, "replaced: int = 37\ndef bar(foo: int) -> int:\n return foo\n\nbiz: int = bar(42)\n", )
def test_replace_sequence_extract(self) -> None: def _reverse_params( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.ensure_type(node, cst.FunctionDef).with_changes( # pyre-ignore We know "params" is a Sequence[Parameters] but asserting that # to pyre is difficult. params=cst.Parameters( params=list(reversed(extraction["params"])))) # Verify that we can still extract sequences with replace. original = cst.parse_module( "def bar(baz: int, foo: int, ) -> int:\n return baz + foo\n") replaced = cst.ensure_type( m.replace( original, m.FunctionDef(params=m.Parameters(params=m.SaveMatchedNode( [m.ZeroOrMore(m.Param())], "params"))), _reverse_params, ), cst.Module, ).code self.assertEqual( replaced, "def bar(foo: int, baz: int, ) -> int:\n return baz + foo\n")
def test_replace_add_one_to_foo_args(self) -> None: def _add_one_to_arg( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return node.deep_replace( # This can be either a node or a sequence, pyre doesn't know. cst.ensure_type(extraction["arg"], cst.CSTNode), # Grab the arg and add one to its value. cst.Integer( str( int( cst.ensure_type(extraction["arg"], cst.Integer).value) + 1)), ) # Verify way more complex transform behavior. original = cst.parse_module( "foo: int = 37\ndef bar(baz: int) -> int:\n return baz\n\nbiz: int = bar(41)\n" ) replaced = cst.ensure_type( m.replace( original, m.Call( func=m.Name("bar"), args=[m.Arg(m.SaveMatchedNode(m.Integer(), "arg"))], ), _add_one_to_arg, ), cst.Module, ).code self.assertEqual( replaced, "foo: int = 37\ndef bar(baz: int) -> int:\n return baz\n\nbiz: int = bar(42)\n", )
def test_replace_simple_sentinel(self) -> None: # Verify behavior when there's a sentinel as a replacement original = cst.parse_module( "def bar(x: int, y: int) -> bool:\n return False\n") replaced = cst.ensure_type( m.replace(original, m.Param(), cst.RemoveFromParent()), cst.Module).code self.assertEqual(replaced, "def bar() -> bool:\n return False\n")
def test_replace_simple(self) -> None: # Verify behavior when there's a static node as a replacement original = cst.parse_module( "foo: bool = True\ndef bar() -> bool:\n return False\n") replaced = cst.ensure_type( m.replace(original, m.Name("True") | m.Name("False"), cst.Name("boolean")), cst.Module, ).code self.assertEqual( replaced, "foo: bool = boolean\ndef bar() -> bool:\n return boolean\n")
def transform_module_impl(self, tree: cst.Module) -> cst.Module: if self.iterations == 0: return tree self.iterations -= 1 return cst.ensure_type( m.replace( tree, m.Integer(), lambda node, _: node.with_changes(value=str(int(node.value) + 1)), ), cst.Module, )
def test_replace_add_one(self) -> None: def _add_one( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.Integer( str(int(cst.ensure_type(node, cst.Integer).value) + 1)) # Verify slightly more complex transform behavior. original = cst.parse_module( "foo: int = 36\ndef bar() -> int:\n return 41\n") replaced = cst.ensure_type(m.replace(original, m.Integer(), _add_one), cst.Module).code self.assertEqual(replaced, "foo: int = 37\ndef bar() -> int:\n return 42\n")
def test_replace_actual(self) -> None: def _swap_bools( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.Name("True" if cst.ensure_type(node, cst.Name).value == "False" else "False") # Verify behavior when there's lots to replace. original = cst.parse_module( "foo: bool = True\ndef bar() -> bool:\n return False\n") replaced = cst.ensure_type( m.replace(original, m.Name("True") | m.Name("False"), _swap_bools), cst.Module, ).code self.assertEqual( replaced, "foo: bool = False\ndef bar() -> bool:\n return True\n")
def test_replace_noop(self) -> None: def _swap_bools( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.Name("True" if cst.ensure_type(node, cst.Name).value == "False" else "False") # Verify behavior when there's nothing to replace. original = cst.parse_module( "foo: int = 5\ndef bar() -> str:\n return 's'\n") replaced = cst.ensure_type( m.replace(original, m.Name("True") | m.Name("False"), _swap_bools), cst.Module, ) # Should be identical tree contents self.assertTrue(original.deep_equals(replaced)) # However, should be a new tree by identity self.assertNotEqual(id(original), id(replaced))
def test_replace_add_one() -> None: original = cst.parse_module( "foo: int = 36\ndef bar() -> int:\n return 41\n") replaced = m.replace(original, m.Integer(), _add_one) assert replaced == "foo: int = 37\ndef bar() -> int:\n return 42\n"