From 2fd28f5e3085558d235ab4a9eaeb1899d6517539 Mon Sep 17 00:00:00 2001
From: Synvoya <16019863+Synvoya@users.noreply.github.com>
Date: Sun, 12 Jul 2026 00:12:20 +1000
Subject: [PATCH] fix(scala): capture trait declarations as class-like nodes

trait_definition was missing from _SCALA_CONFIG.class_types, so Scala
traits produced no node and their heritage (`extends Base with Logging`)
was dropped. Add trait_definition to the Scala class_types; the existing
extends_clause handler then emits the inherits/mixes_in edges.

Adds a regression test for a trait with an extends + with clause.
---
 graphify/extract.py     |  4 +++-
 tests/test_languages.py | 13 +++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/graphify/extract.py b/graphify/extract.py
index 59d4f8283..584df37fe 100644
--- a/graphify/extract.py
+++ b/graphify/extract.py
@@ -780,7 +780,9 @@ def _get_c_func_name(node, source: bytes) -> str | None:
 
 _SCALA_CONFIG = LanguageConfig(
     ts_module="tree_sitter_scala",
-    class_types=frozenset({"class_definition", "object_definition"}),
+    # traits are class-like containers with their own heritage (extends / with),
+    # so they need a node and the heritage walk just like classes and objects.
+    class_types=frozenset({"class_definition", "object_definition", "trait_definition"}),
     function_types=frozenset({"function_definition"}),
     import_types=frozenset({"import_declaration"}),
     call_types=frozenset({"call_expression"}),
diff --git a/tests/test_languages.py b/tests/test_languages.py
index f610fbc40..a3a1406fc 100644
--- a/tests/test_languages.py
+++ b/tests/test_languages.py
@@ -691,6 +691,19 @@ def test_scala_splits_inherits_and_mixes_in():
     assert ("HttpClient", "Loggable") in _edge_labels(r, "mixes_in")
 
 
+def test_scala_trait_definition_heritage(tmp_path):
+    """A `trait` is a class-like container and must get a node plus heritage
+    edges. `trait_definition` was missing from the Scala class_types, so traits
+    produced no node and their `extends`/`with` edges were dropped.
+    """
+    f = tmp_path / "greeter.scala"
+    f.write_text("trait Greeter extends Base with Logging { def greet(): Unit }\n")
+    r = extract_scala(f)
+    assert any("Greeter" in l for l in _labels(r))
+    assert ("Greeter", "Base") in _edge_labels(r, "inherits")
+    assert ("Greeter", "Logging") in _edge_labels(r, "mixes_in")
+
+
 def test_scala_constructor_parameter_field_context():
     r = extract_scala(FIXTURES / "sample.scala")
     assert ("HttpClient", "Config") in _edge_labels(r, "references", "field")