Replacer.py

import ast
import json

class NameReplacer(ast.NodeTransformer):
    def __init__(self, white_list, variables, functions, classes):
        self.white_list = set(white_list)
        self.variables = set(variables)
        self.functions = set(functions)
        self.classes = set(classes)

        self.replacement_map = {}
        self.var_count = 1
        self.func_count = 1
        self.class_count = 1
        
        # Track variable scopes to ensure consistent replacement
        self.scopes = [{}]  # Stack of scopes, innermost scope last
        self.current_scope = self.scopes[-1]

    def get_new_name(self, prefix, count):
        return f"{prefix}{count}"

    def replace_and_remember(self, old_name, prefix, counter_attr):
        if old_name in self.replacement_map:
            return self.replacement_map[old_name]
        else:
            count = getattr(self, counter_attr)
            new_name = self.get_new_name(prefix, count)
            setattr(self, counter_attr, count + 1)
            self.replacement_map[old_name] = new_name
            # Update current scope
            self.current_scope[old_name] = new_name
            return new_name
            
    def get_from_scopes(self, name):
        """Check all scopes for a name, from innermost to outermost"""
        for scope in reversed(self.scopes):
            if name in scope:
                return scope[name]
        return None

    def enter_scope(self):
        """Enter a new scope (e.g., function or class)"""
        self.scopes.append({})
        self.current_scope = self.scopes[-1]
        
    def exit_scope(self):
        """Exit the current scope"""
        if len(self.scopes) > 1:  # Keep at least one scope
            self.scopes.pop()
        self.current_scope = self.scopes[-1]

    def visit_Name(self, node):
        # Check if this name is already in our scopes
        replacement = self.get_from_scopes(node.id)
        if replacement:
            node.id = replacement
        # Handle variable references and class references
        elif node.id in self.variables and node.id not in self.white_list and node.id != "self":
            node.id = self.replace_and_remember(node.id, "var", "var_count")
        elif node.id in self.classes and node.id not in self.white_list:
            node.id = self.replace_and_remember(node.id, "class", "class_count")
        elif node.id in self.functions and node.id not in self.white_list:
            node.id = self.replace_and_remember(node.id, "function", "func_count")
        return self.generic_visit(node)

    def visit_FunctionDef(self, node):
        # Enter a new scope for function body
        self.enter_scope()
        
        # Handle function name
        if node.name in self.functions and node.name not in self.white_list and not node.name.startswith("__"):
            node.name = self.replace_and_remember(node.name, "function", "func_count")
        
        # Handle function parameters
        for arg in node.args.args:
            if arg.arg != "self" and arg.arg in self.variables and arg.arg not in self.white_list:
                arg.arg = self.replace_and_remember(arg.arg, "var", "var_count")
        
        # Process function body
        self.generic_visit(node)
        
        # Exit function scope
        self.exit_scope()
        return node

    def visit_ClassDef(self, node):
        # Enter a new scope for class body
        self.enter_scope()
        
        if node.name in self.classes and node.name not in self.white_list:
            node.name = self.replace_and_remember(node.name, "class", "class_count")
            
        # Process class body
        self.generic_visit(node)
        
        # Exit class scope
        self.exit_scope()
        return node
        
    def visit_Assign(self, node):
        # Process right side first (values being assigned)
        node.value = self.visit(node.value)
        
        # Then process the targets (variables being assigned to)
        for target in node.targets:
            if isinstance(target, ast.Name):
                if target.id in self.variables and target.id not in self.white_list and target.id != "self":
                    # Remember this variable in the current scope
                    target.id = self.replace_and_remember(target.id, "var", "var_count")
            else:
                self.visit(target)
        
        return node

    def visit_Attribute(self, node):
        # Visit the value part first (e.g., the object in obj.attr)
        node.value = self.visit(node.value)
        
        # Handle method calls and attribute access
        if isinstance(node.value, ast.Name) and node.value.id == "self":
            # Class attribute or method reference through self
            if node.attr in self.functions and node.attr not in self.white_list:
                node.attr = self.replace_and_remember(node.attr, "function", "func_count")
            elif node.attr in self.variables and node.attr not in self.white_list:
                node.attr = self.replace_and_remember(node.attr, "var", "var_count")
        elif node.attr in self.functions and node.attr not in self.white_list:
            # Method call on other objects
            node.attr = self.replace_and_remember(node.attr, "function", "func_count")
        elif node.attr in self.variables and node.attr not in self.white_list:
            # Attribute access
            node.attr = self.replace_and_remember(node.attr, "var", "var_count")
            
        return node
    
    def visit_Call(self, node):
        # Handle function calls
        if isinstance(node.func, ast.Name):
            replacement = self.get_from_scopes(node.func.id)
            if replacement:
                node.func.id = replacement
            elif node.func.id in self.functions and node.func.id not in self.white_list:
                node.func.id = self.replace_and_remember(node.func.id, "function", "func_count")
            elif node.func.id in self.classes and node.func.id not in self.white_list:
                # Class instantiation
                node.func.id = self.replace_and_remember(node.func.id, "class", "class_count")
        elif isinstance(node.func, ast.Attribute):
            # Visit the value part first
            node.func.value = self.visit(node.func.value)
            # Method calls like obj.method()
            if node.func.attr in self.functions and node.func.attr not in self.white_list:
                node.func.attr = self.replace_and_remember(node.func.attr, "function", "func_count")
        
        # Handle function arguments
        for i, arg in enumerate(node.args):
            node.args[i] = self.visit(arg)
        
        # Handle keyword arguments
        for keyword in node.keywords:
            if keyword.arg in self.variables and keyword.arg not in self.white_list:
                keyword.arg = self.replace_and_remember(keyword.arg, "var", "var_count")
            keyword.value = self.visit(keyword.value)
        
        return node
    
    def visit_arg(self, node):
        # Handle function arguments in lambda expressions
        if node.arg in self.variables and node.arg not in self.white_list and node.arg != "self":
            node.arg = self.replace_and_remember(node.arg, "var", "var_count")
        return self.generic_visit(node)
    
    def visit_Import(self, node):
        # Don't modify import statements
        return node
    
    def visit_ImportFrom(self, node):
        # Don't modify import statements
        return node
        
    def visit_For(self, node):
        # Process the iterable first
        node.iter = self.visit(node.iter)
        
        # Enter a new scope for the loop
        self.enter_scope()
        
        # Then handle the target (variable being assigned in the loop)
        if isinstance(node.target, ast.Name):
            if node.target.id in self.variables and node.target.id not in self.white_list:
                node.target.id = self.replace_and_remember(node.target.id, "var", "var_count")
        else:
            node.target = self.visit(node.target)
            
        # Process loop body
        for stmt in node.body:
            self.visit(stmt)
            
        # Process else clause if present
        for stmt in node.orelse:
            self.visit(stmt)
            
        # Exit loop scope
        self.exit_scope()
        return node

def safe_parse(code):
    try:
        return ast.parse(code)
    except SyntaxError:
        return None

def replace_code(code, white_list, code_elements):
    tree = safe_parse(code)
    if tree is None:
        print("Cannot parse code using AST, keeping original code")
        return code, {}
        
    try:
        # Use NameReplacer to replace identifiers in the code
        replacer = NameReplacer(white_list, 
                               code_elements['variables'], 
                               code_elements['functions'], 
                               code_elements['classes'])
        new_tree = replacer.visit(tree)
        ast.fix_missing_locations(new_tree)
        new_code = ast.unparse(new_tree)
        replacement_map = {v: k for k, v in replacer.replacement_map.items()}  # Reverse key-value pairs
        return new_code, json.dumps(replacement_map)
    except Exception as e:
        print(f"Error using AST-based replacement: {e}, falling back to original code")
        return code, "{}"