def _replace_or_remove(
parent: cst.CSTNode,
original_node: cst.CSTNode,
replacement_node: Union[cst.CSTNode, cst.RemovalSentinel],
) -> cst.CSTNode:
if isinstance(replacement_node, cst.RemovalSentinel):
return cst.ensure_type(parent.deep_remove(original_node), cst.CSTNode)
else:
return parent.deep_replace(original_node, replacement_node)
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)),
)