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"
