From 54c147ea98b13b69fe33eae2b570237f671b54ed Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:37:58 +0530 Subject: [PATCH] feat(go): basic Go language support (#17) Adds tree-sitter-go and a GO_QUERIES module covering structs, interfaces, functions, methods, struct fields, imports and calls. Go models types as `type_declaration -> type_spec` with the concrete shape on the `type` field, so structs and interfaces are matched on the type_spec rather than on a dedicated node. Methods carry a `receiver` and are a separate node type (method_declaration) from plain functions, so Area() lands as METHOD while describe()/main() land as FUNCTION. The calls query handles both shapes Go uses -- a bare identifier (describe()) and a selector (fmt.Println()) -- capturing the method name for the latter. Query names reuse the existing generic mapping in node_extractor, so no per-language dispatch was needed. Verified on a sample with an interface, struct, method, two functions and both call shapes: INTERFACE Shape / STRUCT Rect / FUNCTION describe, main / METHOD Area FIELD W, H / IMPORTS 1 Call edges additionally require the _extract_call_edges fix from #50; with that applied locally this sample yields exactly the two intra-file edges (describe -> Area, main -> describe) and correctly excludes external calls such as fmt.Println. This PR does not include that fix. test_unsupported_language.py used .go as its example of an unsupported extension, which is no longer true -- switched to .rb. Suite: 3 failed, 188 passed (baseline on main: 3 failed, 174 passed) -- same three pre-existing failures, fixed separately in #51. --- ast_rag/dto/enums.py | 1 + ast_rag/services/parsing/__init__.py | 3 + ast_rag/services/parsing/go.py | 64 ++++++++++ ast_rag/services/parsing/parser_manager.py | 3 + pyproject.toml | 1 + requirements.txt | 1 + tests/test_go_parsing.py | 138 +++++++++++++++++++++ tests/test_unsupported_language.py | 12 +- 8 files changed, 217 insertions(+), 6 deletions(-) create mode 100644 ast_rag/services/parsing/go.py create mode 100644 tests/test_go_parsing.py diff --git a/ast_rag/dto/enums.py b/ast_rag/dto/enums.py index e6e8eed..08ac14c 100644 --- a/ast_rag/dto/enums.py +++ b/ast_rag/dto/enums.py @@ -71,6 +71,7 @@ class Language(str, Enum): PYTHON = "python" TYPESCRIPT = "typescript" TSX = "tsx" + GO = "go" class BlockType(str, Enum): diff --git a/ast_rag/services/parsing/__init__.py b/ast_rag/services/parsing/__init__.py index 3c8be0c..d6913de 100644 --- a/ast_rag/services/parsing/__init__.py +++ b/ast_rag/services/parsing/__init__.py @@ -11,6 +11,7 @@ from ast_rag.services.parsing.rust import RUST_QUERIES from ast_rag.services.parsing.python import PYTHON_QUERIES from ast_rag.services.parsing.typescript import TYPESCRIPT_QUERIES +from ast_rag.services.parsing.go import GO_QUERIES LANGUAGE_QUERIES: dict[str, dict[str, str]] = { "java": JAVA_QUERIES, @@ -18,6 +19,7 @@ "rust": RUST_QUERIES, "python": PYTHON_QUERIES, "typescript": TYPESCRIPT_QUERIES, + "go": GO_QUERIES, # TSX grammar is a superset of TypeScript's (adds JSX nodes), so the # existing TypeScript queries apply unchanged. "tsx": TYPESCRIPT_QUERIES, @@ -35,6 +37,7 @@ "RUST_QUERIES", "PYTHON_QUERIES", "TYPESCRIPT_QUERIES", + "GO_QUERIES", "BlockExtractor", "ParserManager", "NodeExtractor", diff --git a/ast_rag/services/parsing/go.py b/ast_rag/services/parsing/go.py new file mode 100644 index 0000000..bfaa868 --- /dev/null +++ b/ast_rag/services/parsing/go.py @@ -0,0 +1,64 @@ +""" +go.py - Tree-sitter S-expression queries for Go. + +BASIC extraction: structs, interfaces, functions, methods, imports, calls. + +Go models types through `type_declaration -> type_spec`, with the concrete +shape (`struct_type` / `interface_type`) hanging off the `type` field, so +structs and interfaces are matched on the type_spec rather than on a +dedicated node. Methods are `method_declaration` (they carry a `receiver`); +plain functions are `function_declaration`. +""" + +from __future__ import annotations + +GO_QUERIES: dict[str, str] = { + "struct_defs": """ +(type_spec + name: (type_identifier) @name + type: (struct_type) @body +) @node +""", + "interface_defs": """ +(type_spec + name: (type_identifier) @name + type: (interface_type) @body +) @node +""", + "function_defs": """ +(function_declaration + name: (identifier) @name + parameters: (parameter_list) @params +) @node +""", + "method_defs": """ +(method_declaration + receiver: (parameter_list) @receiver + name: (field_identifier) @name + parameters: (parameter_list) @params +) @node +""", + "field_defs": """ +(field_declaration + name: (field_identifier) @name + type: (_) @field_type +) @node +""", + "imports": """ +(import_spec + path: (interpreted_string_literal) @path +) @node +""", + # `callee_name` is what EdgeExtractor._extract_call_edges reads. Go calls + # are either a bare identifier (`helper()`) or a selector + # (`fmt.Println()`); for the latter the method name is the useful half. + "calls": """ +[ + (call_expression + function: (identifier) @callee_name) + (call_expression + function: (selector_expression + field: (field_identifier) @callee_name)) +] @node +""", +} diff --git a/ast_rag/services/parsing/parser_manager.py b/ast_rag/services/parsing/parser_manager.py index 4025d8f..8012d4c 100644 --- a/ast_rag/services/parsing/parser_manager.py +++ b/ast_rag/services/parsing/parser_manager.py @@ -20,6 +20,7 @@ import tree_sitter_cpp as tscpp import tree_sitter_java as tsjava import tree_sitter_rust as tsrust +import tree_sitter_go as tsgo import tree_sitter_python as tspython import tree_sitter_typescript as tsts import tree_sitter as ts @@ -50,6 +51,7 @@ ".h": "cpp", ".java": "java", ".rs": "rust", + ".go": "go", ".py": "python", ".ts": "typescript", ".tsx": "tsx", @@ -139,6 +141,7 @@ def _init_languages(self) -> None: "cpp": ts.Language(tscpp.language()), "java": ts.Language(tsjava.language()), "rust": ts.Language(tsrust.language()), + "go": ts.Language(tsgo.language()), "python": ts.Language(tspython.language()), "typescript": ts.Language(tsts.language_typescript()), "tsx": ts.Language(tsts.language_tsx()), diff --git a/pyproject.toml b/pyproject.toml index 0b60b19..7c8d68d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "tree-sitter-rust>=0.23", "tree-sitter-python>=0.23", "tree-sitter-typescript>=0.23", + "tree-sitter-go>=0.23", # Graph database "neo4j>=5.14", # Vector store (Qdrant — Python 3.14 compatible) diff --git a/requirements.txt b/requirements.txt index 7cab137..d7eb278 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,7 @@ tree-sitter-java>=0.23 tree-sitter-rust>=0.23 tree-sitter-python>=0.23 tree-sitter-typescript>=0.23 +tree-sitter-go>=0.23 # Graph database driver neo4j>=5.14 diff --git a/tests/test_go_parsing.py b/tests/test_go_parsing.py new file mode 100644 index 0000000..ec466c0 --- /dev/null +++ b/tests/test_go_parsing.py @@ -0,0 +1,138 @@ +"""Basic Go extraction (issue #17). + +Go models types as `type_declaration -> type_spec`, with the concrete shape on +the `type` field, so structs and interfaces are matched on the type_spec rather +than on a dedicated node. Methods carry a `receiver` and are a distinct node +type from plain functions. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from tree_sitter import Query, QueryCursor + +from ast_rag.models import NodeKind +from ast_rag.services.parsing import LANGUAGE_QUERIES +from ast_rag.services.parsing.go import GO_QUERIES +from ast_rag.services.parsing.parser_manager import EXT_TO_LANG, ParserManager + +GO_SRC = b""" +package main + +import "fmt" + +type Shape interface { + Area() float64 +} + +type Rect struct { + W float64 + H float64 +} + +func (r Rect) Area() float64 { + return r.W * r.H +} + +func describe(s Shape) string { + return fmt.Sprintf("area=%v", s.Area()) +} + +func main() { + fmt.Println(describe(Rect{W: 2, H: 3})) +} +""" + + +@pytest.fixture(scope="module") +def pm() -> ParserManager: + return ParserManager() + + +@pytest.fixture() +def parsed(pm: ParserManager, tmp_path: Path): + path = tmp_path / "main.go" + path.write_bytes(GO_SRC) + tree = pm.parse_file(str(path), source=GO_SRC) + assert tree is not None, "Go source failed to parse" + nodes = pm.extract_nodes(tree, str(path), "go") + edges = pm.extract_edges(tree, nodes, str(path), "go", source=GO_SRC) + return nodes, edges + + +def _named(nodes, kind: NodeKind): + return {n.name for n in nodes if n.kind == kind} + + +def test_go_is_registered(): + assert EXT_TO_LANG[".go"] == "go" + assert "go" in LANGUAGE_QUERIES + + +def test_language_detected_from_extension(pm: ParserManager, tmp_path: Path): + path = tmp_path / "x.go" + path.write_bytes(b"package main\n") + assert pm.detect_language(str(path)) == "go" + + +def test_structs_and_interfaces_extracted(parsed): + nodes, _ = parsed + assert "Rect" in _named(nodes, NodeKind.STRUCT) + assert "Shape" in _named(nodes, NodeKind.INTERFACE) + + +def test_functions_and_methods_distinguished(parsed): + nodes, _ = parsed + functions = _named(nodes, NodeKind.FUNCTION) + methods = _named(nodes, NodeKind.METHOD) + assert {"describe", "main"} <= functions + # Area has a receiver, so it is a method rather than a function + assert "Area" in methods + assert "Area" not in functions + + +def test_struct_fields_extracted(parsed): + nodes, _ = parsed + assert {"W", "H"} <= _named(nodes, NodeKind.FIELD) + + +def test_imports_extracted(parsed): + _, edges = parsed + kinds = {str(e.kind) for e in edges} + assert any("IMPORTS" in k for k in kinds) + + +def test_calls_query_captures_bare_and_selector_calls(): + """Both `helper()` and `pkg.Helper()` must yield a callee_name. + + Asserted at the query level: turning these matches into CALLS edges also + requires the _extract_call_edges fix, which is a separate change. + """ + import tree_sitter as ts + import tree_sitter_go as tsgo + + lang = ts.Language(tsgo.language()) + tree = ts.Parser(lang).parse(GO_SRC) + matches = list(QueryCursor(Query(lang, GO_QUERIES["calls"])).matches(tree.root_node)) + + names = set() + for _, md in matches: + cap = md.get("callee_name") + if cap is None: + continue + node = cap[0] if isinstance(cap, list) else cap + names.add(node.text.decode()) + + assert "describe" in names, "bare identifier call not captured" + assert "Println" in names, "selector call not captured" + assert "Area" in names + + +@pytest.mark.parametrize("query_name", sorted(GO_QUERIES)) +def test_every_go_query_compiles(query_name: str): + import tree_sitter as ts + import tree_sitter_go as tsgo + + Query(ts.Language(tsgo.language()), GO_QUERIES[query_name]) diff --git a/tests/test_unsupported_language.py b/tests/test_unsupported_language.py index ad97a5f..bd2bd94 100644 --- a/tests/test_unsupported_language.py +++ b/tests/test_unsupported_language.py @@ -65,12 +65,12 @@ def test_format_lists_every_language(self) -> None: class TestParserManagerUnsupported: def test_returns_none_and_warns(self, pm: ParserManager, tmp_path: Path, caplog) -> None: - path = tmp_path / "main.go" - path.write_text("package main\n", encoding="utf-8") + path = tmp_path / "main.rb" + path.write_text("puts 1\n", encoding="utf-8") with caplog.at_level("WARNING"): assert pm.parse_file(str(path)) is None messages = [rec.message for rec in caplog.records] - assert any(".go" in m and "Supported languages" in m for m in messages) + assert any(".rb" in m and "Supported languages" in m for m in messages) def test_extension_less_file_warns(self, pm: ParserManager, tmp_path: Path, caplog) -> None: path = tmp_path / "Makefile" @@ -95,12 +95,12 @@ def test_supported_file_does_not_warn(self, pm: ParserManager, tmp_path: Path, c class TestParsingServiceUnsupported: def test_value_error_lists_supported_languages(self, tmp_path: Path) -> None: service = ParsingService() - path = tmp_path / "main.go" - path.write_text("package main\n", encoding="utf-8") + path = tmp_path / "main.rb" + path.write_text("puts 1\n", encoding="utf-8") with pytest.raises(ValueError) as exc_info: service.parse_file(str(path)) message = str(exc_info.value) - assert "'.go'" in message + assert "'.rb'" in message assert "Supported languages" in message assert "java" in message