diff --git a/graphify/extract.py b/graphify/extract.py index fa103cb..a3f9b29 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -2065,7 +2065,11 @@ _JAVA_CONFIG = LanguageConfig( ts_module="tree_sitter_java", # record_declaration shares class_declaration's name/body/interfaces fields, # so it becomes a first-class type node instead of an isolated file (#1373). - class_types=frozenset({"class_declaration", "interface_declaration", "record_declaration"}), + # Enums and annotation declarations use the same name/body contract. + class_types=frozenset({ + "class_declaration", "interface_declaration", "record_declaration", + "enum_declaration", "annotation_type_declaration", + }), function_types=frozenset({"method_declaration", "constructor_declaration"}), import_types=frozenset({"import_declaration"}), # object_creation_expression (`new Foo(...)`) is handled by a dedicated Java diff --git a/tests/test_languages.py b/tests/test_languages.py index d2fd439..ba41156 100644 --- a/tests/test_languages.py +++ b/tests/test_languages.py @@ -399,6 +399,36 @@ def test_java_type_annotations_have_attribute_context(tmp_path): assert ("CheckoutService", "Entity") in refs +def test_java_enum_and_annotation_declarations_are_type_nodes(tmp_path): + source = tmp_path / "TypeDeclarations.java" + source.write_text( + "enum PaymentStatus { PENDING, PAID }\n" + "@interface Audited {}\n" + "class Order { PaymentStatus status; }\n" + "@Audited class CheckoutService {}\n" + ) + + result = extract_java(source) + + assert ("TypeDeclarations.java", "PaymentStatus") in _edge_labels( + result, "contains" + ) + assert ("TypeDeclarations.java", "Audited") in _edge_labels(result, "contains") + assert ("Order", "PaymentStatus") in _edge_labels( + result, "references", "field" + ) + assert ("CheckoutService", "Audited") in _edge_labels( + result, "references", "attribute" + ) + definitions = { + node["label"]: node + for node in result["nodes"] + if node.get("label") in {"PaymentStatus", "Audited"} + } + assert definitions["PaymentStatus"].get("source_file") == str(source) + assert definitions["Audited"].get("source_file") == str(source) + + def test_csharp_field_type_references_have_field_context(): r = extract_csharp(FIXTURES / "sample.cs") refs = _references(r)