da

transformers/attribute_transformer.py

@@ -1,26 +1,36 @@
+"""
+@file transformers/attribute_transformer.py
+@brief Attribute access transformer for consistent class member renaming.
+@details Ensures references like self.method and self.attr are kept consistent
+         with class and attribute mappings produced by analysis.
+"""
+
 import ast
 from typing import Dict, Optional
 
 class AttributeTransformer(ast.NodeTransformer):
     """
-    Transforms attribute access expressions (obj.attr) consistently,
-    especially for class method calls and attribute accesses.
+    @brief Transform attribute access expressions for consistency.
+    @details Handles self.method and self.attr translations based on provided
+             mappings to keep calls and attributes aligned with obfuscated names.
     """
     
     def __init__(self, class_attr_mapping: Dict[str, Dict[str, str]], class_renames: Dict[str, str]):
         """
-        Initialize with mapping dictionaries.
-        
-        Args:
-            class_attr_mapping: Maps class_name -> {attr_name -> obfuscated_attr_name}
-            class_renames: Maps original_class_name -> obfuscated_class_name
+        @brief Initialize transformer with mapping dictionaries.
+        @param class_attr_mapping Maps class_name -> {attr_name -> obfuscated_attr_name}
+        @param class_renames Maps original_class_name -> obfuscated_class_name
         """
         self.class_attr_mapping = class_attr_mapping
         self.class_renames = class_renames
         self.current_class: Optional[str] = None
     
     def visit_ClassDef(self, node):
-        """Keep track of the current class being processed"""
+        """
+        @brief Track the current class and process its body.
+        @param node ast.ClassDef being visited.
+        @return ast.ClassDef Potentially modified node.
+        """
         old_class = self.current_class
         
         # Get the obfuscated name for this class
@@ -38,7 +48,11 @@ class AttributeTransformer(ast.NodeTransformer):
         return node
     
     def visit_Attribute(self, node):
-        """Transform attribute access like self.method to use consistent names"""
+        """
+        @brief Transform attribute access for consistency within a class.
+        @param node ast.Attribute node to transform.
+        @return ast.AST Updated attribute node.
+        """
         # First process any nested attributes
         node = self.generic_visit(node)