diff-code-change-range 0.0.1__py3-none-any.whl

diff_code_change_range/structure_extractor.py

@@ -0,0 +1,750 @@
+"""Code structure extractor using tree-sitter for Java and Kotlin."""
+
+import sys
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict, Any, Callable
+from enum import Enum
+
+import tree_sitter
+from tree_sitter import Language, Parser
+
+# Try to import language libraries, handle if not available
+try:
+    import tree_sitter_java as ts_java
+    JAVA_LANGUAGE = Language(ts_java.language())
+    JAVA_AVAILABLE = True
+except ImportError:
+    JAVA_AVAILABLE = False
+    JAVA_LANGUAGE = None
+
+try:
+    import tree_sitter_kotlin as ts_kotlin
+    KOTLIN_LANGUAGE = Language(ts_kotlin.language())
+    KOTLIN_AVAILABLE = True
+except ImportError:
+    KOTLIN_AVAILABLE = False
+    KOTLIN_LANGUAGE = None
+
+try:
+    import tree_sitter_python as ts_python
+    PYTHON_LANGUAGE = Language(ts_python.language())
+    PYTHON_AVAILABLE = True
+except ImportError:
+    PYTHON_AVAILABLE = False
+    PYTHON_LANGUAGE = None
+
+
+class NodeType(Enum):
+    """Types of code nodes."""
+    FILE = "file"
+    CLASS = "class"
+    INTERFACE = "interface"
+    OBJECT = "object"
+    ENUM = "enum"
+    FUNCTION = "function"
+    METHOD = "method"
+    MEMBER = "member"
+    VARIABLE = "variable"  # Module-level variables (Python)
+
+
+@dataclass
+class CodeNode:
+    """Represents a node in the code structure tree."""
+    name: str
+    node_type: NodeType
+    line_range: tuple  # (start_line, end_line) - 1-based, inclusive
+    children: List['CodeNode'] = field(default_factory=list)
+    is_affected: bool = False
+
+
+def extract_structure(source_code: str, file_path: str) -> Optional[CodeNode]:
+    """
+    Extract code structure from source code using tree-sitter.
+    
+    Args:
+        source_code: The source code to parse
+        file_path: Path to the file (used to determine language)
+        
+    Returns:
+        CodeNode representing the file structure, or None if parsing fails
+    """
+    if not source_code:
+        return None
+    
+    language = _get_language_for_file(file_path)
+    if not language:
+        print(f"Warning: Unsupported file type: {file_path}", file=sys.stderr)
+        return None
+    
+    parser = Parser(language)
+    
+    try:
+        tree = parser.parse(bytes(source_code, 'utf8'))
+    except Exception as e:
+        print(f"Warning: Failed to parse {file_path}: {e}", file=sys.stderr)
+        return None
+    
+    root_node = tree.root_node
+    
+    # Extract file-level nodes based on language
+    if file_path.endswith('.java'):
+        children = _extract_java_file_nodes(root_node, source_code)
+    elif file_path.endswith('.kt'):
+        children = _extract_kotlin_file_nodes(root_node, source_code)
+    elif file_path.endswith('.py'):
+        children = _extract_python_file_nodes(root_node, source_code)
+    else:
+        return None
+    
+    # Calculate file line range
+    if root_node.child_count > 0:
+        start_line = root_node.start_point[0] + 1  # Convert to 1-based
+        end_line = root_node.end_point[0] + 1
+    else:
+        start_line = 1
+        end_line = source_code.count('\n') + 1
+    
+    return CodeNode(
+        name=file_path,
+        node_type=NodeType.FILE,
+        line_range=(start_line, end_line),
+        children=children
+    )
+
+
+def _get_language_for_file(file_path: str) -> Optional[Language]:
+    """Get the appropriate Language for a file based on extension."""
+    if file_path.endswith('.java') and JAVA_AVAILABLE:
+        return JAVA_LANGUAGE
+    elif file_path.endswith('.kt') and KOTLIN_AVAILABLE:
+        return KOTLIN_LANGUAGE
+    elif file_path.endswith('.py') and PYTHON_AVAILABLE:
+        return PYTHON_LANGUAGE
+    return None
+
+
+def _extract_java_file_nodes(root_node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract top-level nodes from Java source."""
+    nodes = []
+    
+    # Cursor to traverse the tree
+    cursor = root_node.walk()
+    
+    # Visit all children of root
+    if cursor.goto_first_child():
+        while True:
+            node = cursor.node
+            node_type = node.type
+            
+            if node_type == 'class_declaration':
+                class_node = _extract_java_class(node, source_code)
+                if class_node:
+                    nodes.append(class_node)
+            elif node_type == 'interface_declaration':
+                interface_node = _extract_java_interface(node, source_code)
+                if interface_node:
+                    nodes.append(interface_node)
+            elif node_type == 'enum_declaration':
+                enum_node = _extract_java_enum(node, source_code)
+                if enum_node:
+                    nodes.append(enum_node)
+            elif node_type == 'field_declaration':
+                # Top-level fields (in class context, but handle here for completeness)
+                pass  # Fields are handled within class extraction
+            
+            if not cursor.goto_next_sibling():
+                break
+    
+    return nodes
+
+
+def _extract_java_class(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Java class node."""
+    name = _extract_node_name(node, source_code)
+    children = _extract_java_class_children(node, source_code)
+    
+    return CodeNode(
+        name=name,
+        node_type=NodeType.CLASS,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1),
+        children=children
+    )
+
+
+def _extract_java_interface(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Java interface node."""
+    name = _extract_node_name(node, source_code)
+    children = _extract_java_class_children(node, source_code)  # Similar structure
+    
+    return CodeNode(
+        name=name,
+        node_type=NodeType.INTERFACE,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1),
+        children=children
+    )
+
+
+def _extract_java_enum(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Java enum node."""
+    name = _extract_node_name(node, source_code)
+    children = _extract_java_class_children(node, source_code)
+    
+    return CodeNode(
+        name=name,
+        node_type=NodeType.ENUM,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1),
+        children=children
+    )
+
+
+def _extract_java_class_children(class_node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract children (methods, fields) from a Java class."""
+    children = []
+    
+    # Find the class body
+    body_node = None
+    for child in class_node.children:
+        if child.type == 'class_body':
+            body_node = child
+            break
+    
+    if not body_node:
+        return children
+    
+    for child in body_node.children:
+        node_type = child.type
+        
+        if node_type == 'method_declaration':
+            method = _extract_java_method(child, source_code)
+            if method:
+                children.append(method)
+        elif node_type == 'field_declaration':
+            fields = _extract_java_fields(child, source_code)
+            children.extend(fields)
+        elif node_type == 'constructor_declaration':
+            method = _extract_java_constructor(child, source_code)
+            if method:
+                children.append(method)
+        elif node_type in ('class_declaration', 'interface_declaration', 'enum_declaration'):
+            # Inner classes
+            if node_type == 'class_declaration':
+                inner = _extract_java_class(child, source_code)
+            elif node_type == 'interface_declaration':
+                inner = _extract_java_interface(child, source_code)
+            else:
+                inner = _extract_java_enum(child, source_code)
+            if inner:
+                children.append(inner)
+    
+    return children
+
+
+def _extract_java_method(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Java method node."""
+    name = _extract_node_name(node, source_code)
+    return CodeNode(
+        name=name,
+        node_type=NodeType.METHOD,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+    )
+
+
+def _extract_java_constructor(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Java constructor as a method node."""
+    name = _extract_node_name(node, source_code)
+    return CodeNode(
+        name=name,
+        node_type=NodeType.METHOD,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+    )
+
+
+def _extract_java_fields(node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract field declarations from a Java field_declaration node."""
+    fields = []
+    
+    # field_declaration contains type and declarator(s)
+    for child in node.children:
+        if child.type == 'variable_declarator':
+            field_name = _extract_node_name(child, source_code)
+            fields.append(CodeNode(
+                name=field_name,
+                node_type=NodeType.MEMBER,
+                line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+            ))
+    
+    return fields
+
+
+def _extract_kotlin_file_nodes(root_node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract top-level nodes from Kotlin source."""
+    nodes = []
+    
+    cursor = root_node.walk()
+    
+    if cursor.goto_first_child():
+        while True:
+            node = cursor.node
+            node_type = node.type
+            
+            if node_type == 'class_declaration':
+                class_node = _extract_kotlin_class(node, source_code)
+                if class_node:
+                    nodes.append(class_node)
+            elif node_type == 'object_declaration':
+                obj_node = _extract_kotlin_object(node, source_code)
+                if obj_node:
+                    nodes.append(obj_node)
+            elif node_type == 'function_declaration':
+                func_node = _extract_kotlin_function(node, source_code)
+                if func_node:
+                    nodes.append(func_node)
+            elif node_type == 'property_declaration':
+                prop_nodes = _extract_kotlin_properties(node, source_code)
+                nodes.extend(prop_nodes)
+            
+            if not cursor.goto_next_sibling():
+                break
+    
+    return nodes
+
+
+def _extract_kotlin_class(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Kotlin class node."""
+    name = _extract_node_name(node, source_code)
+    children = _extract_kotlin_class_children(node, source_code)
+    
+    # Check if it's an enum class
+    node_type = NodeType.CLASS
+    for child in node.children:
+        if child.type == 'enum_class_body':
+            node_type = NodeType.ENUM
+            break
+    
+    return CodeNode(
+        name=name,
+        node_type=node_type,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1),
+        children=children
+    )
+
+
+def _extract_kotlin_object(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Kotlin object node."""
+    name = _extract_node_name(node, source_code)
+    children = _extract_kotlin_class_children(node, source_code)
+    
+    return CodeNode(
+        name=name,
+        node_type=NodeType.OBJECT,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1),
+        children=children
+    )
+
+
+def _extract_kotlin_function(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Kotlin function node (top-level or class-level)."""
+    name = _extract_node_name(node, source_code)
+    return CodeNode(
+        name=name,
+        node_type=NodeType.FUNCTION,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+    )
+
+
+def _extract_kotlin_properties(node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract property declarations from Kotlin."""
+    properties = []
+    
+    # Find variable declarations within property_declaration
+    for child in node.children:
+        if child.type == 'variable_declaration':
+            prop_name = _extract_node_name(child, source_code)
+            properties.append(CodeNode(
+                name=prop_name,
+                node_type=NodeType.MEMBER,
+                line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+            ))
+    
+    # If no variable_declaration found, try to get name from property_declaration itself
+    if not properties:
+        name = _extract_node_name(node, source_code)
+        if name:
+            properties.append(CodeNode(
+                name=name,
+                node_type=NodeType.MEMBER,
+                line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+            ))
+    
+    return properties
+
+
+def _extract_kotlin_class_children(class_node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract children from a Kotlin class or object."""
+    children = []
+    
+    # Find the class body
+    body_node = None
+    for child in class_node.children:
+        if child.type in ('class_body', 'enum_class_body', 'object_body'):
+            body_node = child
+            break
+    
+    if not body_node:
+        return children
+    
+    for child in body_node.children:
+        node_type = child.type
+        
+        if node_type == 'function_declaration':
+            # Determine if it's a method (inside class) or function (top-level)
+            func = _extract_kotlin_function(child, source_code)
+            if func:
+                func.node_type = NodeType.METHOD  # Inside class, it's a method
+                children.append(func)
+        elif node_type == 'property_declaration':
+            props = _extract_kotlin_properties(child, source_code)
+            children.extend(props)
+        elif node_type in ('class_declaration', 'object_declaration'):
+            if node_type == 'class_declaration':
+                inner = _extract_kotlin_class(child, source_code)
+            else:
+                inner = _extract_kotlin_object(child, source_code)
+            if inner:
+                children.append(inner)
+    
+    return children
+
+
+def _extract_node_name(node: tree_sitter.Node, source_code: str) -> str:
+    """Extract the name/identifier from a node."""
+    # Look for identifier child
+    for child in node.children:
+        if child.type == 'identifier':
+            return source_code[child.start_byte:child.end_byte]
+    
+    # For some nodes, the name might be in a different field
+    # Try to find any identifier recursively in immediate children
+    for child in node.children[:3]:  # Look in first few children
+        if child.type == 'identifier':
+            return source_code[child.start_byte:child.end_byte]
+    
+    return "<anonymous>"
+
+
+# ============================================================================
+# Python Extraction Functions
+# ============================================================================
+
+def _extract_python_file_nodes(root_node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract top-level nodes from Python source."""
+    nodes = []
+    
+    cursor = root_node.walk()
+    
+    if cursor.goto_first_child():
+        while True:
+            node = cursor.node
+            node_type = node.type
+            
+            if node_type == 'class_definition':
+                class_node = _extract_python_class(node, source_code)
+                if class_node:
+                    nodes.append(class_node)
+            elif node_type == 'function_definition':
+                func_node = _extract_python_function(node, source_code)
+                if func_node:
+                    nodes.append(func_node)
+            elif node_type == 'expression_statement':
+                # Check for assignment (module-level variable)
+                assign_node = _extract_python_module_variable(node, source_code)
+                if assign_node:
+                    nodes.append(assign_node)
+            elif node_type == 'assignment':
+                # Direct assignment at module level
+                var_node = _extract_python_variable_from_assignment(node, source_code)
+                if var_node:
+                    nodes.append(var_node)
+            
+            if not cursor.goto_next_sibling():
+                break
+    
+    return nodes
+
+
+def _extract_python_class(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Python class node."""
+    name = _extract_node_name(node, source_code)
+    children = _extract_python_class_children(node, source_code)
+    
+    return CodeNode(
+        name=name,
+        node_type=NodeType.CLASS,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1),
+        children=children
+    )
+
+
+def _extract_python_function(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Python function node (module-level)."""
+    name = _build_python_function_name(node, source_code)
+    return CodeNode(
+        name=name,
+        node_type=NodeType.FUNCTION,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+    )
+
+
+def _build_python_function_name(node: tree_sitter.Node, source_code: str) -> str:
+    """Extract function name with async prefix if applicable."""
+    # Check if it's an async function by looking at first child
+    is_async = False
+    for child in node.children:
+        if child.type == 'async':
+            is_async = True
+            break
+        elif child.type == 'def':
+            break
+    
+    name = _extract_node_name(node, source_code)
+    if is_async:
+        return f"async {name}"
+    return name
+
+
+def _extract_python_module_variable(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract module-level variable from expression_statement containing assignment."""
+    # expression_statement may contain assignment
+    for child in node.children:
+        if child.type == 'assignment':
+            return _extract_python_variable_from_assignment(child, source_code)
+    return None
+
+
+def _extract_python_variable_from_assignment(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a variable from an assignment node."""
+    # Get the left-hand side (pattern/identifier)
+    for child in node.children:
+        if child.type in ('identifier', 'pattern'):
+            name = source_code[child.start_byte:child.end_byte]
+            return CodeNode(
+                name=name,
+                node_type=NodeType.VARIABLE,
+                line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+            )
+        elif child.type == 'pattern_list' or child.type == 'tuple_pattern':
+            # Handle tuple unpacking: a, b = 1, 2
+            # Return the first variable as primary
+            for subchild in child.children:
+                if subchild.type == 'identifier':
+                    name = source_code[subchild.start_byte:subchild.end_byte]
+                    return CodeNode(
+                        name=name,
+                        node_type=NodeType.VARIABLE,
+                        line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+                    )
+    return None
+
+
+def _extract_python_class_children(class_node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract children (methods, class variables, nested classes) from a Python class."""
+    children = []
+    
+    # Find the class body (block)
+    body_node = None
+    for child in class_node.children:
+        if child.type == 'block':
+            body_node = child
+            break
+    
+    if not body_node:
+        return children
+    
+    for child in body_node.children:
+        node_type = child.type
+        
+        if node_type == 'function_definition':
+            method = _extract_python_method(child, source_code)
+            if method:
+                children.append(method)
+        elif node_type == 'decorated_definition':
+            # Handle decorated functions/methods
+            method = _extract_python_decorated_method(child, source_code)
+            if method:
+                children.append(method)
+        elif node_type == 'class_definition':
+            # Handle nested classes
+            inner_class = _extract_python_class(child, source_code)
+            if inner_class:
+                children.append(inner_class)
+        elif node_type == 'expression_statement':
+            # Check for class variable assignment
+            member_node = _extract_python_class_variable(child, source_code)
+            if member_node:
+                children.append(member_node)
+        elif node_type == 'assignment':
+            # Direct assignment in class body
+            member = _extract_python_variable_from_assignment(child, source_code)
+            if member:
+                member.node_type = NodeType.MEMBER
+                children.append(member)
+    
+    return children
+
+
+def _extract_python_decorated_method(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a decorated Python method node."""
+    # Find the function_definition inside decorated_definition
+    for child in node.children:
+        if child.type == 'function_definition':
+            return _extract_python_method(node, source_code)
+    return None
+
+
+def _extract_python_method(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract a Python method node."""
+    name = _build_decorated_name(node, source_code)
+    
+    # Find the actual function definition node for instance variable extraction
+    func_node = node
+    if node.type == 'decorated_definition':
+        for child in node.children:
+            if child.type == 'function_definition':
+                func_node = child
+                break
+    
+    # Also extract instance variables from __init__ and other methods
+    instance_vars = _extract_instance_variables(func_node, source_code)
+    
+    method_node = CodeNode(
+        name=name,
+        node_type=NodeType.METHOD,
+        line_range=(node.start_point[0] + 1, node.end_point[0] + 1),
+        children=instance_vars
+    )
+    
+    return method_node
+
+
+def _build_decorated_name(node: tree_sitter.Node, source_code: str) -> str:
+    """Build method name with decorator prefixes."""
+    # Collect decorators and check for async
+    decorators = []
+    is_async = False
+    func_node = node
+    
+    # If this is a decorated_definition, look inside for the function
+    if node.type == 'decorated_definition':
+        for child in node.children:
+            if child.type == 'function_definition':
+                func_node = child
+                break
+        # Extract decorators from decorated_definition
+        for child in node.children:
+            if child.type == 'decorator':
+                dec_name = _extract_decorator_name(child, source_code)
+                if dec_name:
+                    decorators.append(dec_name)
+    
+    # Check for async in function_definition
+    for child in func_node.children:
+        if child.type == 'async':
+            is_async = True
+            break
+    
+    # Get base name from function node
+    base_name = _extract_node_name(func_node, source_code)
+    
+    # Build full name
+    name_parts = []
+    if decorators:
+        name_parts.append('@' + ',@'.join(decorators))
+    if is_async:
+        name_parts.append('async')
+    name_parts.append(base_name)
+    
+    return ' '.join(name_parts)
+
+
+def _extract_decorator_name(decorator_node: tree_sitter.Node, source_code: str) -> str:
+    """Extract decorator name from a decorator node."""
+    # decorator contains identifier or call
+    for child in decorator_node.children:
+        if child.type == 'identifier':
+            return source_code[child.start_byte:child.end_byte]
+        elif child.type == 'call':
+            # Decorator with arguments: @decorator(...)
+            for subchild in child.children:
+                if subchild.type == 'identifier':
+                    return source_code[subchild.start_byte:subchild.end_byte]
+        elif child.type == 'attribute':
+            # Decorator like @abc.classmethod
+            return source_code[child.start_byte:child.end_byte]
+    return ""
+
+
+def _extract_python_class_variable(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract class-level variable from expression_statement."""
+    for child in node.children:
+        if child.type == 'assignment':
+            member = _extract_python_variable_from_assignment(child, source_code)
+            if member:
+                member.node_type = NodeType.MEMBER
+                return member
+    return None
+
+
+def _extract_instance_variables(node: tree_sitter.Node, source_code: str) -> List[CodeNode]:
+    """Extract instance variables (self.x) from method body."""
+    instance_vars = []
+    
+    # Find the block (method body)
+    body_node = None
+    for child in node.children:
+        if child.type == 'block':
+            body_node = child
+            break
+    
+    if not body_node:
+        return instance_vars
+    
+    # Look for expression_statement containing assignment
+    for child in body_node.children:
+        if child.type == 'expression_statement':
+            for subchild in child.children:
+                if subchild.type == 'assignment':
+                    var = _extract_self_assignment(subchild, source_code)
+                    if var:
+                        instance_vars.append(var)
+        elif child.type == 'assignment':
+            var = _extract_self_assignment(child, source_code)
+            if var:
+                instance_vars.append(var)
+    
+    return instance_vars
+
+
+def _extract_self_assignment(node: tree_sitter.Node, source_code: str) -> Optional[CodeNode]:
+    """Extract instance variable from self.x = ... assignment."""
+    # Check if left side is attribute (self.x)
+    for child in node.children:
+        if child.type == 'attribute':
+            # Check if it's self.something
+            is_self = False
+            attr_name = ""
+            
+            for subchild in child.children:
+                if subchild.type == 'identifier' and source_code[subchild.start_byte:subchild.end_byte] == 'self':
+                    is_self = True
+                elif subchild.type == 'identifier' and is_self:
+                    attr_name = source_code[subchild.start_byte:subchild.end_byte]
+                elif subchild.type == 'type':  # The . operator
+                    pass
+            
+            if is_self and attr_name:
+                return CodeNode(
+                    name=attr_name,
+                    node_type=NodeType.MEMBER,
+                    line_range=(node.start_point[0] + 1, node.end_point[0] + 1)
+                )
+    return None