Source code for cosapp.core.ir.visitors

import ast

from .ops import Node




class NodeVisitor(ast.NodeVisitor):
    """Visit nodes."""



    def generic_visit(self, node):
        for _, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, (ast.AST, Node)):
                        self.visit(item)
            elif isinstance(value, (ast.AST, Node)):
                self.visit(value)






class NodeTransformer(ast.NodeTransformer):
    """Visit and transform nodes."""



    def generic_visit(self, node):
        if isinstance(node, list):
            return [self.visit(value) for value in node]

        for field_, old_value in ast.iter_fields(node):
            if isinstance(old_value, list):
                new_values = []
                for value in old_value:
                    if isinstance(value, (ast.AST, Node)):
                        value = self.visit(value)
                        if value is None:
                            continue
                        elif not isinstance(value, (ast.AST, Node)):
                            new_values.extend(value)
                            continue
                    new_values.append(value)
                old_value[:] = new_values
            elif isinstance(old_value, (ast.AST, Node)):
                new_node = self.visit(old_value)
                if new_node is None:
                    delattr(node, field_)
                else:
                    setattr(node, field_, new_node)
        return node