codegraphcontext-rust 0.5.5__cp311-cp311-win_amd64.whl

codegraphcontext/tools/languages/go.py

@@ -0,0 +1,499 @@
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple
+from codegraphcontext.utils.debug_log import debug_log, info_logger, error_logger, warning_logger
+from codegraphcontext.utils.tree_sitter_manager import execute_query
+
+
+GO_QUERIES = {
+    "functions": """
+        (function_declaration
+            name: (identifier) @name
+            parameters: (parameter_list) @params
+        ) @function_node
+        
+        (method_declaration
+            receiver: (parameter_list) @receiver
+            name: (field_identifier) @name
+            parameters: (parameter_list) @params
+        ) @function_node
+    """,
+    "structs": """
+        (type_declaration
+            (type_spec
+                name: (type_identifier) @name
+                type: (struct_type) @struct_body
+            )
+        ) @struct_node
+    """,
+    "interfaces": """
+        (type_declaration
+            (type_spec
+                name: (type_identifier) @name
+                type: (interface_type) @interface_body
+            )
+        ) @interface_node
+    """,
+    "imports": """
+        (import_spec
+            path: (interpreted_string_literal) @path
+        )
+    """,
+    "calls": """
+        (call_expression
+            function: (identifier) @name
+        )
+        (call_expression
+            function: (selector_expression
+                field: (field_identifier) @name
+            )
+        )
+    """,
+    "variables": """
+        (var_declaration
+            (var_spec
+                name: (identifier) @name
+            )
+        )
+        (short_var_declaration
+            left: (expression_list
+                (identifier) @name
+            )
+        )
+    """,
+}
+
+class GoTreeSitterParser:
+    """A Go-specific parser using tree-sitter, encapsulating language-specific logic."""
+
+    def __init__(self, generic_parser_wrapper):
+        self.generic_parser_wrapper = generic_parser_wrapper
+        self.language_name = generic_parser_wrapper.language_name
+        self.language = generic_parser_wrapper.language
+        self.parser = generic_parser_wrapper.parser
+        self.index_source = False
+
+    def _get_node_text(self, node) -> str:
+        return node.text.decode('utf-8')
+
+    def _get_parent_context(self, node, types=('function_declaration', 'method_declaration', 'type_declaration')):
+        curr = node.parent
+        while curr:
+            if curr.type in types:
+                if curr.type == 'type_declaration':
+                    type_spec = curr.child_by_field_name('type_spec')
+                    if type_spec:
+                        name_node = type_spec.child_by_field_name('name')
+                        return self._get_node_text(name_node) if name_node else None, curr.type, curr.start_point[0] + 1
+                else:
+                    name_node = curr.child_by_field_name('name')
+                    return self._get_node_text(name_node) if name_node else None, curr.type, curr.start_point[0] + 1
+            curr = curr.parent
+        return None, None, None
+
+    def _calculate_complexity(self, node):
+        """
+        Compute a simple cyclomatic complexity score from the Go AST.
+
+        We treat each decision/control-flow construct as +1:
+        - if/for/switch/select
+        - switch cases (case_clause) and select clauses (comm_clause)
+        - logical operators (&&, ||) inside binary_expression as +1
+        """
+        # Note: tree-sitter-go node types differ from other languages and from
+        # what we'd typically expect (e.g. it's `switch`/`case` rather than
+        # `switch_statement`/`case_clause`).
+        decision_node_types = {
+            # top-level control constructs
+            "if_statement",
+            "for_statement",
+            "switch",
+            "select_statement",
+            "select",
+            # switch variants (tree-sitter grammar)
+            "expression_switch_statement",
+            "type_switch_statement",
+            # switch case branches
+            "case",
+            "expression_case",
+            "default_case",
+            # select communication branches
+            "communication_case",
+        }
+
+        count = 1
+
+        def traverse(n):
+            nonlocal count
+            if n.type in decision_node_types:
+                count += 1
+                # Still traverse children because nested constructs also contribute.
+            elif n.type == "binary_expression":
+                # Only count logical operators, not all binary expressions/comparisons.
+                # Example patterns: "a && b", "a || b".
+                try:
+                    txt = self._get_node_text(n)
+                except Exception:
+                    txt = ""
+                if "&&" in txt or "||" in txt:
+                    count += 1
+
+            for child in n.children:
+                traverse(child)
+
+        traverse(node)
+        return count
+
+    def _get_docstring(self, func_node):
+        """Extract Go doc comment preceding the function."""
+        prev_sibling = func_node.prev_sibling
+        while prev_sibling and prev_sibling.type in ('comment', '\n', ' '):
+            if prev_sibling.type == 'comment':
+                comment_text = self._get_node_text(prev_sibling)
+                if comment_text.startswith('//'):
+                    return comment_text.strip()
+            prev_sibling = prev_sibling.prev_sibling
+        return None
+
+    def parse(self, path: Path, is_dependency: bool = False, index_source: bool = False) -> Dict:
+        """Parses a file and returns its structure in a standardized dictionary format."""
+        # This method orchestrates the parsing of a single file.
+        # It calls specialized `_find_*` methods for each language construct.
+        # The returned dictionary should map a specific key (e.g., 'functions', 'interfaces')
+        # to a list of dictionaries, where each dictionary represents a single code construct.
+        # The GraphBuilder will then use these keys to create nodes with corresponding labels.
+        self.index_source = index_source
+        with open(path, "r", encoding="utf-8") as f:
+            source_code = f.read()
+
+        tree = self.parser.parse(bytes(source_code, "utf8"))
+        root_node = tree.root_node
+
+        functions = self._find_functions(root_node)
+        structs = self._find_structs(root_node)
+        interfaces = self._find_interfaces(root_node)
+        imports = self._find_imports(root_node)
+        function_calls = self._find_calls(root_node)
+        variables = self._find_variables(root_node)
+
+        return {
+            "path": str(path),
+            "functions": functions,
+            "classes": structs,
+            "interfaces": interfaces,
+            "variables": variables,
+            "imports": imports,
+            "function_calls": function_calls,
+            "is_dependency": is_dependency,
+            "lang": self.language_name,
+        }
+
+    def _find_functions(self, root_node):
+        functions = []
+        query_str = GO_QUERIES['functions']
+
+        captures_by_function = {}
+
+        for node, capture_name in execute_query(self.language, query_str, root_node):
+            if capture_name == 'function_node':
+                func_id = node.id
+                if func_id not in captures_by_function:
+                    captures_by_function[func_id] = {
+                        'node': node,
+                        'name': None,
+                        'params': None,
+                        'receiver': None
+                    }
+            elif capture_name == 'name':
+                func_node = self._find_function_node_for_name(node)
+                if func_node:
+                    func_id = func_node.id
+                    if func_id not in captures_by_function:
+                        captures_by_function[func_id] = {
+                            'node': func_node,
+                            'name': None,
+                            'params': None,
+                            'receiver': None
+                        }
+                    captures_by_function[func_id]['name'] = self._get_node_text(node)
+            elif capture_name == 'params':
+                func_node = self._find_function_node_for_params(node)
+                if func_node:
+                    func_id = func_node.id
+                    if func_id not in captures_by_function:
+                        captures_by_function[func_id] = {
+                            'node': func_node,
+                            'name': None,
+                            'params': None,
+                            'receiver': None
+                        }
+                    captures_by_function[func_id]['params'] = node
+            elif capture_name == 'receiver':
+                func_node = node.parent
+                if func_node and func_node.type == 'method_declaration':
+                    func_id = func_node.id
+                    if func_id not in captures_by_function:
+                        captures_by_function[func_id] = {
+                            'node': func_node,
+                            'name': None,
+                            'params': None,
+                            'receiver': None
+                        }
+                    captures_by_function[func_id]['receiver'] = node
+
+        for func_id, data in captures_by_function.items():
+            if data['name']:
+                func_node = data['node']
+                name = data['name']
+
+                args = []
+                if data['params']:
+                    args = self._extract_parameters(data['params'])
+
+                receiver_type = None
+                if data['receiver']:
+                    receiver_type = self._extract_receiver(data['receiver'])
+
+                context, context_type, context_line = self._get_parent_context(func_node)
+                class_context = receiver_type or (context if context_type == 'type_declaration' else None)
+
+                docstring = self._get_docstring(func_node)
+
+                func_data = {
+                    "name": name,
+                    "line_number": func_node.start_point[0] + 1,
+                    "end_line": func_node.end_point[0] + 1,
+                    "args": args,
+                    "class_context": class_context,
+                    "decorators": [],
+                    "lang": self.language_name,
+                    "is_dependency": False,
+                    "cyclomatic_complexity": self._calculate_complexity(func_node),
+                }
+                
+                if self.index_source:
+                    func_data["source"] = self._get_node_text(func_node)
+                    func_data["docstring"] = docstring
+                    
+                functions.append(func_data)
+
+        return functions
+
+    def _find_function_node_for_name(self, name_node):
+        current = name_node.parent
+        while current:
+            if current.type in ('function_declaration', 'method_declaration'):
+                return current
+            current = current.parent
+        return None
+
+    def _find_function_node_for_params(self, params_node):
+        current = params_node.parent
+        while current:
+            if current.type in ('function_declaration', 'method_declaration'):
+                return current
+            current = current.parent
+        return None
+
+    def _extract_parameters(self, params_node):
+        params = []
+        if params_node.type == 'parameter_list':
+            for child in params_node.children:
+                if child.type == 'parameter_declaration':
+                    # Handle multiple names for same type: func(x, y int)
+                    # We iterate children and find all identifiers that are not the type node.
+                    type_node = child.child_by_field_name('type')
+                    for grandchild in child.children:
+                        if grandchild.type == 'identifier':
+                            if grandchild.id != (type_node.id if type_node else None):
+                                params.append(self._get_node_text(grandchild))
+                elif child.type == 'variadic_parameter_declaration':
+                    name_node = child.child_by_field_name('name')
+                    if name_node:
+                        params.append(f"...{self._get_node_text(name_node)}")
+        return params
+
+    def _extract_receiver(self, receiver_node):
+        if receiver_node.type == 'parameter_list' and receiver_node.named_child_count > 0:
+            param = receiver_node.named_child(0)
+            type_node = param.child_by_field_name('type')
+            if type_node:
+                type_text = self._get_node_text(type_node)
+                return type_text.strip('*')
+        return None
+
+    def _find_structs(self, root_node):
+        structs = []
+        struct_query_str = GO_QUERIES['structs']
+        for node, capture_name in execute_query(self.language, struct_query_str, root_node):
+            if capture_name == 'name':
+                struct_node = self._find_type_declaration_for_name(node)
+                if struct_node:
+                    name = self._get_node_text(node)
+                    class_data = {
+                        "name": name,
+                        "line_number": struct_node.start_point[0] + 1,
+                        "end_line": struct_node.end_point[0] + 1,
+                        "bases": [],
+                        "decorators": [],
+                        "lang": self.language_name,
+                        "is_dependency": False,
+                    }
+                    if self.index_source:
+                        class_data["source"] = self._get_node_text(struct_node)
+                        class_data["docstring"] = self._get_docstring(struct_node)
+
+                    structs.append(class_data)
+        return structs
+
+    def _find_interfaces(self, root_node):
+        interfaces = []
+        interface_query_str = GO_QUERIES['interfaces']
+        for node, capture_name in execute_query(self.language, interface_query_str, root_node):
+            if capture_name == 'name':
+                interface_node = self._find_type_declaration_for_name(node)
+                if interface_node:
+                    name = self._get_node_text(node)
+                    class_data = {
+                        "name": name,
+                        "line_number": interface_node.start_point[0] + 1,
+                        "end_line": interface_node.end_point[0] + 1,
+                        "bases": [],
+                        "decorators": [],
+                        "lang": self.language_name,
+                        "is_dependency": False,
+                    }
+                    if self.index_source:
+                        class_data["source"] = self._get_node_text(interface_node)
+                        class_data["docstring"] = self._get_docstring(interface_node)
+                        
+                    interfaces.append(class_data)
+        return interfaces
+
+    def _find_type_declaration_for_name(self, name_node):
+        current = name_node.parent
+        while current:
+            if current.type == 'type_declaration':
+                return current
+            current = current.parent
+        return None
+
+    def _find_imports(self, root_node):
+        imports = []
+        query_str = GO_QUERIES['imports']
+        
+        for node, capture_name in execute_query(self.language, query_str, root_node):
+            line_number = node.start_point[0] + 1
+            
+            if capture_name == 'path':
+                path_text = self._get_node_text(node).strip('"')
+                package_name = path_text.split('/')[-1]
+                
+                alias = None
+                import_spec = node.parent
+                if import_spec and import_spec.type == 'import_spec':
+                    alias_node = import_spec.child_by_field_name('name')
+                    if alias_node:
+                        alias = self._get_node_text(alias_node)
+                
+                imports.append({
+                    'name': package_name,
+                    'source': path_text,
+                    'alias': alias,
+                    'line_number': line_number,
+                    'lang': self.language_name
+                })
+
+        return imports
+
+    def _find_calls(self, root_node):
+        calls = []
+        query_str = GO_QUERIES['calls']
+        
+        seen_calls = set()
+
+        for node, capture_name in execute_query(self.language, query_str, root_node):
+            if capture_name == 'name':
+                call_node = node.parent
+                while call_node and call_node.type != 'call_expression':
+                    call_node = call_node.parent
+                
+                if call_node:
+                    name = self._get_node_text(node)
+                    line_number = node.start_point[0] + 1
+                    
+                    call_key = f"{name}_{line_number}"
+                    if call_key in seen_calls:
+                        continue
+                    seen_calls.add(call_key)
+                    
+                    full_name = self._get_node_text(call_node.child_by_field_name('function')) if call_node.child_by_field_name('function') else name
+                    
+                    # Resolve context
+                    context_name, context_type, context_line = self._get_parent_context(node)
+                    
+                    # In Go, methods are defined on types (structs/interfaces). If we are in a method, the context is the method name.
+                    # Ideally we might want the receiver type as "class_context", but this requires more complex AST traversal up to the method declaration's receiver.
+                    # For now, we reuse the context resolution logic.
+                    class_context = None 
+                    
+                    call_data = {
+                        "name": name,
+                        "full_name": full_name,
+                        "line_number": line_number,
+                        "args": [],
+                        "inferred_obj_type": None,
+                        "context": (context_name, context_type, context_line),
+                        "class_context": class_context,
+                        "lang": self.language_name,
+                        "is_dependency": False,
+                    }
+                    calls.append(call_data)
+        
+        return calls
+
+    def _find_variables(self, root_node):
+        variables = []
+        query_str = GO_QUERIES['variables']
+        
+        for node, capture_name in execute_query(self.language, query_str, root_node):
+            if capture_name == 'name':
+                name = self._get_node_text(node)
+                
+                variable_data = {
+                    "name": name,
+                    "line_number": node.start_point[0] + 1,
+                    "value": None,
+                    "type": None,
+                    "context": None,
+                    "class_context": None,
+                    "lang": self.language_name,
+                    "is_dependency": False,
+                }
+                variables.append(variable_data)
+        
+        return variables
+
+def pre_scan_go(files: list[Path], parser_wrapper) -> dict:
+    """Scans Go files to create a map of function/struct names to their file paths."""
+    imports_map = {}
+    query_str = """
+        (function_declaration name: (identifier) @name)
+        (method_declaration name: (field_identifier) @name)
+        (type_declaration (type_spec name: (type_identifier) @name))
+    """
+    
+
+    for path in files:
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                tree = parser_wrapper.parser.parse(bytes(f.read(), "utf8"))
+
+            for capture, _ in execute_query(parser_wrapper.language, query_str, tree.root_node):
+                name = capture.text.decode('utf-8')
+                if name not in imports_map:
+                    imports_map[name] = []
+                imports_map[name].append(str(path.resolve()))
+        except Exception as e:
+            warning_logger(f"Tree-sitter pre-scan failed for {path}: {e}")
+    
+    return imports_map