From 00dd978d17d9916f21fe770b169b7b1aa659f04e Mon Sep 17 00:00:00 2001 From: ivanzhilovich Date: Wed, 8 Jul 2026 22:49:01 +0100 Subject: [PATCH] fix(extract): emit Kotlin enum entries as nodes with case_of edges (#1700 Kotlin half, #1738) Kotlin enum entries weren't extracted: the walker never descended into the enum body (`enum_class_body` wasn't in _KOTLIN_CONFIG.body_fallback_child_types, so _find_body returned None). Add `enum_class_body` to the fallback body types and a `_kotlin_extra_walk` (dispatched for tree_sitter_kotlin, mirroring the Java/Swift handling) that emits each enum_entry as a node with a `case_of` edge to the enum. Re-applied from PR #1738 (@ivanzhl) onto the post-#1737 module layout: the walk engine now lives in graphify/extractors/engine.py while _KOTLIN_CONFIG stays in extract.py. Closes the Kotlin half of #1700 (Java was #1719). Co-Authored-By: Claude Opus 4.8 (1M context) --- graphify/extract.py | 2 +- graphify/extractors/engine.py | 33 +++++++++++++++++++++++++++++++++ tests/fixtures/sample.kt | 6 ++++++ tests/test_languages.py | 8 ++++++++ 4 files changed, 48 insertions(+), 1 deletion(-) diff --git a/graphify/extract.py b/graphify/extract.py index 0e074ab..ffbe47e 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -772,7 +772,7 @@ _KOTLIN_CONFIG = LanguageConfig( # older forks use `simple_identifier`. Accept both so the extractor # works across grammar generations. name_fallback_child_types=("simple_identifier", "identifier"), - body_fallback_child_types=("function_body", "class_body"), + body_fallback_child_types=("function_body", "class_body", "enum_class_body"), function_boundary_types=frozenset({"function_declaration"}), import_handler=_import_kotlin, ) diff --git a/graphify/extractors/engine.py b/graphify/extractors/engine.py index cef89d4..80a9537 100644 --- a/graphify/extractors/engine.py +++ b/graphify/extractors/engine.py @@ -1768,6 +1768,33 @@ def _java_extra_walk(node, source: bytes, file_nid: str, stem: str, str_path: st return True return False + +def _kotlin_extra_walk(node, source: bytes, file_nid: str, stem: str, str_path: str, + nodes: list, edges: list, seen_ids: set, function_bodies: list, + parent_class_nid: str | None, add_node_fn, add_edge_fn, + walk_fn) -> bool: + """Handle enum_entry for Kotlin. Returns True if handled (#1700 Kotlin half).""" + if node.type == "enum_entry" and parent_class_nid: + name_node = None + for child in node.children: + if child.type in ("simple_identifier", "identifier"): + name_node = child + break + if name_node is None: + return True + const_name = _read_text(name_node, source) + line = node.start_point[0] + 1 + const_nid = _make_id(parent_class_nid, const_name) + add_node_fn(const_nid, const_name, line) + add_edge_fn(parent_class_nid, const_nid, "case_of", line) + for child in node.children: + if child.type == "class_body": + for member in child.children: + walk_fn(member, parent_class_nid=const_nid) + return True + return False + + def _read_csharp_type_name(node, source: bytes) -> tuple[str, bool, str] | None: """Resolve a C# type name, whether it was qualified, and its qualifier prefix.""" if node is None: @@ -3299,6 +3326,12 @@ def _extract_generic( parent_class_nid, add_node, add_edge, walk): return + if config.ts_module == "tree_sitter_kotlin": + if _kotlin_extra_walk(node, source, file_nid, stem, str_path, + nodes, edges, seen_ids, function_bodies, + parent_class_nid, add_node, add_edge, walk): + return + if config.ts_module == "tree_sitter_ruby": if _ruby_extra_walk(node, source, file_nid, stem, str_path, nodes, edges, seen_ids, function_bodies, diff --git a/tests/fixtures/sample.kt b/tests/fixtures/sample.kt index db3f281..07f6269 100644 --- a/tests/fixtures/sample.kt +++ b/tests/fixtures/sample.kt @@ -41,3 +41,9 @@ fun createClient(baseUrl: String): HttpClient { val config = Config(baseUrl, 30) return HttpClient(config) } + +enum class ChatType { + NORMAL, + GROUP, + SYSTEM +} diff --git a/tests/test_languages.py b/tests/test_languages.py index 565bf0a..f610fbc 100644 --- a/tests/test_languages.py +++ b/tests/test_languages.py @@ -615,6 +615,14 @@ def test_kotlin_finds_function(): r = extract_kotlin(FIXTURES / "sample.kt") assert any("createClient" in l for l in _labels(r)) +def test_kotlin_enum_entries_have_case_of_edge(): + # #1700 (Kotlin half): enum entries must be nodes with case_of edges to the enum. + r = extract_kotlin(FIXTURES / "sample.kt") + labels = _labels(r) + assert "NORMAL" in labels and "GROUP" in labels and "SYSTEM" in labels + assert ("ChatType", "NORMAL") in _edge_labels(r, "case_of") + assert ("ChatType", "SYSTEM") in _edge_labels(r, "case_of") + def test_kotlin_emits_in_file_calls(): """Regression test for the call-walker `simple_identifier` / `identifier` rename — see graphify-kmp's PythonParityTest."""