def replace_subscripts(vyper_module: vy_ast.Module) -> int:
"""
Find and evaluate Subscript nodes within the Vyper AST, replacing them with
Constant nodes where possible.
Arguments
---------
vyper_module : Module
Top-level Vyper AST node.
Returns
-------
int
Number of nodes that were replaced.
"""
changed_nodes = 0
for node in vyper_module.get_descendants(vy_ast.Subscript, reverse=True):
try:
new_node = node.evaluate()
except UnfoldableNode:
continue
changed_nodes += 1
vyper_module.replace_in_tree(node, new_node)
return changed_nodes
def replace_literal_ops(vyper_module: vy_ast.Module) -> int:
"""
Find and evaluate operation and comparison nodes within the Vyper AST,
replacing them with Constant nodes where possible.
Arguments
---------
vyper_module : Module
Top-level Vyper AST node.
Returns
-------
int
Number of nodes that were replaced.
"""
changed_nodes = 0
node_types = (vy_ast.BoolOp, vy_ast.BinOp, vy_ast.UnaryOp, vy_ast.Compare)
for node in vyper_module.get_descendants(node_types, reverse=True):
try:
new_node = node.evaluate()
except UnfoldableNode:
continue
changed_nodes += 1
vyper_module.replace_in_tree(node, new_node)
return changed_nodes
def replace_builtin_functions(vyper_module: vy_ast.Module) -> int:
"""
Find and evaluate builtin function calls within the Vyper AST, replacing
them with Constant nodes where possible.
Arguments
---------
vyper_module : Module
Top-level Vyper AST node.
Returns
-------
int
Number of nodes that were replaced.
"""
changed_nodes = 0
for node in vyper_module.get_descendants(vy_ast.Call, reverse=True):
if not isinstance(node.func, vy_ast.Name):
continue
name = node.func.id
func = DISPATCH_TABLE.get(name)
if func is None or not hasattr(func, "evaluate"):
continue
try:
new_node = func.evaluate(node) # type: ignore
except UnfoldableNode:
continue
changed_nodes += 1
vyper_module.replace_in_tree(node, new_node)
return changed_nodes
def replace_constant(
vyper_module: vy_ast.Module,
id_: str,
replacement_node: Union[vy_ast.Constant, vy_ast.List],
raise_on_error: bool,
) -> int:
"""
Replace references to a variable name with a literal value.
Arguments
---------
vyper_module : Module
Module-level ast node to perform replacement in.
id_ : str
String representing the `.id` attribute of the node(s) to be replaced.
replacement_node : Constant | List
Vyper ast node representing the literal value to be substituted in.
raise_on_error: bool
Boolean indicating if `UnfoldableNode` exception should be raised or ignored.
Returns
-------
int
Number of nodes that were replaced.
"""
changed_nodes = 0
for node in vyper_module.get_descendants(vy_ast.Name, {"id": id_}, reverse=True):
parent = node.get_ancestor()
if isinstance(parent, vy_ast.Attribute):
# do not replace attributes
continue
if isinstance(parent, vy_ast.Call) and node == parent.func:
# do not replace calls
continue
# do not replace dictionary keys
if isinstance(parent, vy_ast.Dict) and node in parent.keys:
continue
if not isinstance(parent, vy_ast.Index):
# do not replace left-hand side of assignments
assign = node.get_ancestor(
(vy_ast.Assign, vy_ast.AnnAssign, vy_ast.AugAssign)
)
if assign and node in assign.target.get_descendants(include_self=True):
continue
try:
new_node = _replace(node, replacement_node)
except UnfoldableNode:
if raise_on_error:
raise
continue
changed_nodes += 1
vyper_module.replace_in_tree(node, new_node)
return changed_nodes
def replace_constant(
vyper_ast_node: vy_ast.Module,
id_: str,
replacement_node: Union[vy_ast.Constant, vy_ast.List],
) -> None:
"""
Replace references to a variable name with a literal value.
Arguments
---------
vyper_ast_node : Module
Module-level ast node to perform replacement in.
id_ : str
String representing the `.id` attribute of the node(s) to be replaced.
replacement_node : Constant | List
Vyper ast node representing the literal value to be substituted in.
"""
for node in vyper_ast_node.get_descendants(vy_ast.Name, {'id': id_},
reverse=True):
# do not replace attributes or calls
if isinstance(node.get_ancestor(), (vy_ast.Attribute, vy_ast.Call)):
continue
# do not replace dictionary keys
if isinstance(node.get_ancestor(),
vy_ast.Dict) and node in node.get_ancestor().keys:
continue
if not isinstance(node.get_ancestor(), vy_ast.Index):
# do not replace left-hand side of assignments
parent = node.get_ancestor(
(vy_ast.Assign, vy_ast.AnnAssign, vy_ast.AugAssign))
if parent and node in parent.target.get_descendants(
include_self=True):
continue
new_node = _replace(node, replacement_node)
vyper_ast_node.replace_in_tree(node, new_node)
def replace_constant(
vyper_module: vy_ast.Module,
id_: str,
replacement_node: Union[vy_ast.Constant, vy_ast.List, vy_ast.Call],
raise_on_error: bool,
type_: BaseTypeDefinition = None,
) -> int:
"""
Replace references to a variable name with a literal value.
Arguments
---------
vyper_module : Module
Module-level ast node to perform replacement in.
id_ : str
String representing the `.id` attribute of the node(s) to be replaced.
replacement_node : Constant | List | Call
Vyper ast node representing the literal value to be substituted in.
`Call` nodes are for struct constants.
raise_on_error: bool
Boolean indicating if `UnfoldableNode` exception should be raised or ignored.
type_ : BaseTypeDefinition, optional
Type definition to be propagated to type checker.
Returns
-------
int
Number of nodes that were replaced.
"""
is_struct = False
if isinstance(replacement_node, vy_ast.Call) and len(replacement_node.args) == 1:
if isinstance(replacement_node.args[0], vy_ast.Dict):
is_struct = True
changed_nodes = 0
for node in vyper_module.get_descendants(vy_ast.Name, {"id": id_}, reverse=True):
parent = node.get_ancestor()
if isinstance(parent, vy_ast.Call) and node == parent.func:
# do not replace calls that are not structs
if not is_struct:
continue
# do not replace dictionary keys
if isinstance(parent, vy_ast.Dict) and node in parent.keys:
continue
if not node.get_ancestor(vy_ast.Index):
# do not replace left-hand side of assignments
assign = node.get_ancestor((vy_ast.Assign, vy_ast.AnnAssign, vy_ast.AugAssign))
if assign and node in assign.target.get_descendants(include_self=True):
continue
try:
# Codegen may mutate AST (particularly structs).
replacement_node = copy.deepcopy(replacement_node)
new_node = _replace(node, replacement_node, type_=type_)
except UnfoldableNode:
if raise_on_error:
raise
continue
changed_nodes += 1
vyper_module.replace_in_tree(node, new_node)
return changed_nodes