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)), )