Tree Manipulation

def dump_tree(node, name=None, initial_indent='', indentation='    ', 
              maxline=120, maxmerged=80):
    """
    Dump AST in pretty-printed format with indentation.
    
    Parameters:
    - node: AST node to dump
    - name: str, optional name for the node
    - initial_indent: str, initial indentation string (default: '')
    - indentation: str, indentation string per level (default: 4 spaces)
    - maxline: int, maximum line length (default: 120)
    - maxmerged: int, maximum merged line length (default: 80)
    
    Returns:
    str: String representation of the AST
    """

import ast
import astor

code = """
def factorial(n):
    if n <= 1:
        return 1
    return n * factorial(n - 1)
"""

tree = ast.parse(code)
print(astor.dump_tree(tree))

# With custom formatting
formatted_dump = astor.dump_tree(tree, 
                                 name="factorial_function",
                                 initial_indent=">> ",
                                 indentation="  ",
                                 maxline=80)
print(formatted_dump)

def iter_node(node, name='', unknown=None):
    """
    Iterate over AST node attributes or list items.
    
    Parameters:
    - node: AST node or list to iterate over
    - name: str, name prefix for list items (default: '')
    - unknown: set, set to collect unknown attributes (default: None)
    
    Returns:
    Iterator yielding (value, name) pairs
    """

import ast
import astor

code = "x = [1, 2, 3]"
tree = ast.parse(code)

# Iterate over all nodes and their attributes
for node in ast.walk(tree):
    print(f"Node: {type(node).__name__}")
    for value, name in astor.iter_node(node):
        print(f"  {name}: {value}")

def strip_tree(node):
    """
    Strip AST by removing all attributes not in _fields.
    
    Parameters:
    - node: AST node to strip
    
    Returns:
    set: Set of names of all stripped attributes
    """

import ast
import astor

code = "def func(): pass"
tree = ast.parse(code)

# Add some custom attributes
for node in ast.walk(tree):
    node.custom_attr = "custom_value"
    node.debug_info = {"line": 1}

# Strip non-standard attributes
stripped_attrs = astor.strip_tree(tree)
print(f"Stripped attributes: {stripped_attrs}")

class TreeWalk:
    """
    Tree walker that can traverse AST recursively or non-recursively.
    
    Uses MetaFlatten metaclass for advanced traversal patterns.
    """
    
    def __init__(self, node=None):
        """
        Initialize walker and optionally start walking.
        
        Parameters:
        - node: AST node to start walking from (optional)
        """
    
    def setup(self):
        """Set up node-specific handlers."""
    
    def walk(self, node, name=''):
        """
        Walk the tree starting at given node.
        
        Parameters:
        - node: AST node to start walking from
        - name: str, name for the starting node (default: '')
        """
    
    def replace(self, new_node):
        """
        Replace current node with new node.
        
        Parameters:
        - new_node: AST node to replace current node with
        """
    
    @property
    def parent(self):
        """
        Return parent node of current node.
        
        Returns:
        AST node that is the parent of the current node
        """
    
    @property  
    def parent_name(self):
        """
        Return parent node and name tuple.
        
        Returns:
        tuple: (parent_node, name) pair
        """

import ast
import astor

class VariableRenamer(astor.TreeWalk):
    def __init__(self, old_name, new_name):
        self.old_name = old_name
        self.new_name = new_name
        super().__init__()
    
    def visit_Name(self, node):
        if node.id == self.old_name:
            node.id = self.new_name

code = "def func(x): return x + 1"
tree = ast.parse(code)

# Rename variable 'x' to 'value'
renamer = VariableRenamer('x', 'value')
renamer.walk(tree)

print(astor.to_source(tree))

class ExplicitNodeVisitor(ast.NodeVisitor):
    """
    NodeVisitor that removes implicit visits and requires explicit handlers.
    
    Inherits from ast.NodeVisitor but raises errors for unhandled node types.
    """
    
    def visit(self, node):
        """
        Visit a node with explicit handler requirement.
        
        Parameters:
        - node: AST node to visit
        
        Raises:
        AttributeError: If no explicit handler exists for the node type
        """
    
    def abort_visit(self, node):
        """
        Default handler that raises AttributeError for unhandled nodes.
        
        Parameters:
        - node: AST node that has no explicit handler
        
        Raises:
        AttributeError: Always, to indicate missing handler
        """

def allow_ast_comparison():
    """
    Monkey-patch AST nodes to enable equality comparison.
    
    Modifies AST classes in-place to support == and != operations.
    """

def fast_compare(tree1, tree2):
    """
    Optimized comparison of two AST trees for equality.
    
    Parameters:
    - tree1: First AST tree to compare  
    - tree2: Second AST tree to compare
    
    Returns:
    bool: True if trees are structurally equal, False otherwise
    """

import ast
import astor

# Enable AST comparison
astor.allow_ast_comparison()

code1 = "x = 1"
code2 = "x = 1"
code3 = "y = 1"

tree1 = ast.parse(code1)
tree2 = ast.parse(code2)
tree3 = ast.parse(code3)

# Now AST nodes support equality comparison
print(tree1 == tree2)  # True
print(tree1 == tree3)  # False

# Fast comparison for performance-critical code
print(astor.fast_compare(tree1, tree2))  # True
print(astor.fast_compare(tree1, tree3))  # False