fix(extract): emit Java annotation type references (#2426)

Java annotation extraction read only the annotation's name field, dropping
class-literal arguments (@Repeatable(RubricsFor.class), @Uses({A.class,B.class}))
and annotation-member return types (RubricFor[] value()), so a container
annotation became a disconnected island. Emit references edges for both, with
qualified-identity preservation; string/enum annotation arguments are not
treated as type refs, and JDK annotations resolve to sourceless stubs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
oleksii-tumanov
2026-08-17 18:33:35 +01:00
committed by safishamsi
co-authored by Claude Opus 4.8
parent 4be19990aa
commit eadef76a68
3 changed files with 373 additions and 20 deletions
+103 -19
View File
@@ -350,6 +350,7 @@ def _java_collect_type_refs(
generic: bool,
out: list[tuple[str, str]],
skip: frozenset[str] | None = None,
preserve_qualified: bool = False,
) -> None:
"""Walk a Java type expression; append (name, role) tuples."""
if node is None:
@@ -365,18 +366,26 @@ def _java_collect_type_refs(
out.append((name, "generic_arg" if generic else "type"))
return
if t == "scoped_type_identifier":
text = _read_text(node, source).rsplit(".", 1)[-1]
if text and text not in _JAVA_BUILTIN_TYPES:
raw = _read_text(node, source)
simple = raw.rsplit(".", 1)[-1]
text = raw if preserve_qualified else raw.rsplit(".", 1)[-1]
if text and simple not in _JAVA_BUILTIN_TYPES:
out.append((text, "generic_arg" if generic else "type"))
return
if t == "generic_type":
for c in node.children:
if c.type in ("type_identifier", "scoped_type_identifier"):
text = _read_text(c, source).rsplit(".", 1)[-1]
raw = _read_text(c, source)
simple = raw.rsplit(".", 1)[-1]
text = (
raw
if preserve_qualified and c.type == "scoped_type_identifier"
else simple
)
if (
text
and text not in _JAVA_BUILTIN_TYPES
and (c.type == "scoped_type_identifier" or text not in skip)
and simple not in _JAVA_BUILTIN_TYPES
and (c.type == "scoped_type_identifier" or simple not in skip)
):
out.append((text, "generic_arg" if generic else "type"))
break
@@ -384,17 +393,23 @@ def _java_collect_type_refs(
if c.type == "type_arguments":
for arg in c.children:
if arg.is_named:
_java_collect_type_refs(arg, source, True, out, skip)
_java_collect_type_refs(
arg, source, True, out, skip, preserve_qualified
)
return
if t == "array_type":
for c in node.children:
if c.is_named:
_java_collect_type_refs(c, source, generic, out, skip)
_java_collect_type_refs(
c, source, generic, out, skip, preserve_qualified
)
return
if node.is_named:
for c in node.children:
if c.is_named:
_java_collect_type_refs(c, source, generic, out, skip)
_java_collect_type_refs(
c, source, generic, out, skip, preserve_qualified
)
def _java_receiver_type_name(type_node, source: bytes) -> str | None:
@@ -548,21 +563,28 @@ def _java_method_receiver_types(
return table
def _java_annotation_names(declaration_node, source: bytes) -> list[tuple[str, str]]:
"""Collect ``(simple, raw)`` annotation names from a Java declaration's
`modifiers` child. ``raw`` keeps the dotted qualifier of an inline-qualified
annotation (``@org.pkg.Foo``); it equals ``simple`` when unqualified."""
names: list[tuple[str, str]] = []
def _java_annotation_nodes(declaration_node) -> list:
"""Return annotations from a Java declaration's `modifiers` child."""
modifiers = None
for child in declaration_node.children:
if child.type == "modifiers":
modifiers = child
break
if modifiers is None:
return names
for anno in modifiers.children:
if anno.type not in ("marker_annotation", "annotation"):
continue
return []
return [
child
for child in modifiers.children
if child.type in ("marker_annotation", "annotation")
]
def _java_annotation_names(declaration_node, source: bytes) -> list[tuple[str, str]]:
"""Collect ``(simple, raw)`` annotation names from a Java declaration's
`modifiers` child. ``raw`` keeps the dotted qualifier of an inline-qualified
annotation (``@org.pkg.Foo``); it equals ``simple`` when unqualified."""
names: list[tuple[str, str]] = []
for anno in _java_annotation_nodes(declaration_node):
name_node = anno.child_by_field_name("name")
if name_node is None:
for sub in anno.children:
@@ -576,6 +598,35 @@ def _java_annotation_names(declaration_node, source: bytes) -> list[tuple[str, s
names.append((text, raw))
return names
def _java_annotation_class_literal_refs(
declaration_node,
source: bytes,
) -> list[str]:
"""Collect Java type names used as class literals in annotation arguments."""
names: list[str] = []
for anno in _java_annotation_nodes(declaration_node):
arguments = anno.child_by_field_name("arguments")
if arguments is None:
continue
stack = [arguments]
while stack:
current = stack.pop()
if current.type == "class_literal":
type_node = next(
(child for child in current.children if child.is_named),
None,
)
refs: list[tuple[str, str]] = []
_java_collect_type_refs(
type_node, source, False, refs, preserve_qualified=True
)
names.extend(name for name, _role in refs)
continue
stack.extend(child for child in current.children if child.is_named)
return names
def _php_name_text(node, source: bytes) -> str | None:
"""Return the unqualified name text from a PHP `name`/`qualified_name` node."""
if node is None:
@@ -3351,6 +3402,7 @@ def _extract_generic(
if tid.is_named:
_emit_java_parent_type(tid, "inherits", line)
annotation_targets: set[str] = set()
for anno_name, anno_raw in _java_annotation_names(node, source):
# An inline-qualified annotation (`@org.pkg.Foo`) keeps its
# full dotted name so a bare same-named local class can't
@@ -3360,9 +3412,16 @@ def _extract_generic(
if "." in anno_raw and config.ts_module == "tree_sitter_java":
anno_name = anno_raw
target_nid = ensure_named_node(anno_name, line)
if target_nid != class_nid:
if target_nid != class_nid and target_nid not in annotation_targets:
add_edge(class_nid, target_nid, "references", line,
context="attribute")
annotation_targets.add(target_nid)
for ref_name in _java_annotation_class_literal_refs(node, source):
target_nid = ensure_named_node(ref_name, line)
if target_nid != class_nid and target_nid not in annotation_targets:
add_edge(class_nid, target_nid, "references", line,
context="attribute")
annotation_targets.add(target_nid)
if t == "record_declaration":
components = node.child_by_field_name("parameters")
@@ -3668,6 +3727,23 @@ def _extract_generic(
line, context=ctx)
return
if (config.ts_module == "tree_sitter_java"
and t == "annotation_type_element_declaration"
and parent_class_nid):
type_node = node.child_by_field_name("type")
line = node.start_point[0] + 1
refs: list[tuple[str, str]] = []
_java_collect_type_refs(
type_node, source, False, refs, preserve_qualified=True
)
for ref_name, role in refs:
ctx = "generic_arg" if role == "generic_arg" else "return_type"
target_nid = ensure_named_node(ref_name, line)
if target_nid != parent_class_nid:
add_edge(parent_class_nid, target_nid, "references",
line, context=ctx)
return
if (config.ts_module == "tree_sitter_php"
and t == "property_declaration"
and parent_class_nid):
@@ -4018,13 +4094,21 @@ def _extract_generic(
target_nid = ensure_named_node(ref_name, line)
if target_nid != func_nid:
add_edge(func_nid, target_nid, "references", line, context=ctx)
annotation_targets: set[str] = set()
for anno_name, anno_raw in _java_annotation_names(node, source):
# Inline-qualified: keep the dotted name (#2504); see the
# class-level annotation handling above.
target_nid = ensure_named_node(
anno_raw if "." in anno_raw else anno_name, line)
if target_nid != func_nid:
if target_nid != func_nid and target_nid not in annotation_targets:
add_edge(func_nid, target_nid, "references", line, context="attribute")
annotation_targets.add(target_nid)
for ref_name in _java_annotation_class_literal_refs(node, source):
target_nid = ensure_named_node(ref_name, line)
if target_nid != func_nid and target_nid not in annotation_targets:
add_edge(func_nid, target_nid, "references", line,
context="attribute")
annotation_targets.add(target_nid)
if config.ts_module == "tree_sitter_php":
params_container = None
+65 -1
View File
@@ -2491,6 +2491,7 @@ def _resolve_go_type_references(
if new_nodes:
all_nodes.extend(new_nodes)
if not repointed_from:
return
referenced = {endpoint for edge in all_edges
@@ -2574,8 +2575,31 @@ def _resolve_java_type_references(
pkg_by_file[s] = pkg
imports_by_file[s] = imps
# FQN (package.Class) -> definition node id, for type-like defs with a source.
# FQN (package.Class or package.Outer.Inner) -> definition node id, for
# type-like defs with a source. Nested declarations need their containing
# type path because qualified annotation references preserve it.
fqn_to_id: dict[str, str] = {}
node_by_id = {
node.get("id"): node for node in all_nodes if node.get("id")
}
type_parent_by_id: dict[str, str] = {}
for edge in all_edges:
if edge.get("relation") != "contains":
continue
child = node_by_id.get(edge.get("target"))
parent = node_by_id.get(edge.get("source"))
if not child or not parent:
continue
child_src = child.get("source_file", "")
parent_label = parent.get("label", "")
if (
child_src
and parent.get("source_file") == child_src
and parent_label[:1].isupper()
and not parent_label.endswith(".java")
):
type_parent_by_id[child["id"]] = parent["id"]
for node in all_nodes:
label = node.get("label", "")
src = node.get("source_file", "")
@@ -2586,6 +2610,20 @@ def _resolve_java_type_references(
continue
pkg = pkg_by_file[src]
fqn_to_id.setdefault(f"{pkg}.{label}" if pkg else label, nid)
type_path = [label]
seen = {nid}
parent_id = type_parent_by_id.get(nid)
while parent_id and parent_id not in seen:
seen.add(parent_id)
parent = node_by_id[parent_id]
type_path.append(parent["label"])
parent_id = type_parent_by_id.get(parent_id)
if len(type_path) > 1:
nested_name = ".".join(reversed(type_path))
fqn_to_id.setdefault(
f"{pkg}.{nested_name}" if pkg else nested_name,
nid,
)
# Shadow stubs: no source_file, type-like label. Dotted labels are included
# for qualified inline annotations (`@com.example.anno.Loggable`), which the
@@ -2676,6 +2714,32 @@ def _resolve_java_type_references(
if new_nodes:
all_nodes.extend(new_nodes)
# Bare imported and inline-qualified annotation references can start on
# different stubs, then converge on one source-backed node above. Collapse
# only indistinguishable Java attribute-reference facts after that rewire.
seen_attribute_refs: set[tuple] = set()
deduped_edges: list[dict] = []
for edge in all_edges:
if (
edge.get("relation") == "references"
and edge.get("context") == "attribute"
and edge.get("source_file", "") in pkg_by_file
):
key = (
edge.get("source"),
edge.get("target"),
edge.get("relation"),
edge.get("context"),
edge.get("source_file"),
edge.get("source_location"),
)
if key in seen_attribute_refs:
continue
seen_attribute_refs.add(key)
deduped_edges.append(edge)
all_edges[:] = deduped_edges
if not repointed_from:
return
+205
View File
@@ -580,6 +580,211 @@ def test_java_type_annotations_have_attribute_context(tmp_path):
assert ("CheckoutService", "Entity") in refs
def test_java_annotation_class_literal_arguments(tmp_path):
source = tmp_path / "AnnotationArguments.java"
source.write_text(
"class PrimaryRule {}\n"
"class BackupRule {}\n"
"@interface UsesRules { Class<?>[] value(); }\n"
"@UsesRules({PrimaryRule.class, BackupRule.class, java.lang.String.class})\n"
"class RuleHandler {\n"
" @UsesRules(PrimaryRule.class)\n"
" void apply() {}\n"
"}\n"
)
result = extract_java(source)
refs = _edge_labels(result, "references", "attribute")
assert ("RuleHandler", "PrimaryRule") in refs
assert ("RuleHandler", "BackupRule") in refs
assert ("apply", "PrimaryRule") in refs
assert ("RuleHandler", "java.lang.String") not in refs
def test_java_annotation_references_are_not_duplicated(tmp_path):
from graphify.extract import extract
annotation = tmp_path / "pkg" / "Uses.java"
annotation.parent.mkdir()
annotation.write_text(
"package pkg;\n"
"public @interface Uses { Class<?>[] value(); }\n"
)
consumer = tmp_path / "app" / "Consumer.java"
consumer.parent.mkdir()
consumer.write_text(
"package app;\n"
"import pkg.Uses;\n"
"@Uses({pkg.Uses.class, pkg.Uses.class})\n"
"class Consumer {\n"
" @Uses({pkg.Uses.class, pkg.Uses.class})\n"
" void apply() {}\n"
"}\n"
)
result = extract([annotation, consumer], cache_root=tmp_path / "cache")
refs = [
pair
for pair in _edge_labels(result, "references", "attribute")
if pair in {("Consumer", "Uses"), ("apply", "Uses")}
]
assert refs.count(("Consumer", "Uses")) == 1
assert refs.count(("apply", "Uses")) == 1
def test_java_annotation_class_literal_keeps_qualified_type_identity(tmp_path):
from graphify.extract import extract
internal = tmp_path / "internal" / "Rule.java"
internal.parent.mkdir()
internal.write_text("package internal;\npublic class Rule {}\n")
outer = tmp_path / "internal" / "Outer.java"
outer.write_text(
"package internal;\n"
"public class Outer { public static class Nested {} }\n"
)
consumer = tmp_path / "app" / "Consumer.java"
consumer.parent.mkdir()
consumer.write_text(
"package app;\n"
"@interface Uses { Class<?>[] value(); }\n"
"@Uses({internal.Rule.class, internal.Outer.Nested.class, "
"external.Rule.class})\n"
"public class Consumer {}\n"
)
result = extract([internal, outer, consumer], cache_root=tmp_path / "cache")
by_id = {node["id"]: node for node in result["nodes"]}
targets = [
by_id[edge["target"]]
for edge in result["edges"]
if edge.get("relation") == "references"
and edge.get("context") == "attribute"
and edge.get("source_file", "").endswith("Consumer.java")
and by_id[edge["target"]].get("label")
in {"Rule", "Nested", "external.Rule"}
]
assert any(
target.get("label") == "Rule"
and target.get("source_file", "").endswith("internal/Rule.java")
for target in targets
)
assert any(
target.get("label") == "external.Rule" and not target.get("source_file")
for target in targets
)
assert any(
target.get("label") == "Nested"
and target.get("source_file", "").endswith("Outer.java")
for target in targets
)
def test_java_repeatable_annotation_references_container_and_element_type(tmp_path):
from graphify.extract import extract
annotation = tmp_path / "RubricFor.java"
annotation.write_text(
"package com.example;\n"
"import java.lang.annotation.Repeatable;\n"
"@Repeatable(RubricsFor.class)\n"
"public @interface RubricFor {}\n"
)
container = tmp_path / "RubricsFor.java"
container.write_text(
"package com.example;\n"
"public @interface RubricsFor {\n"
" RubricFor[] value();\n"
"}\n"
)
result = extract([annotation, container], cache_root=tmp_path / "cache")
assert ("RubricFor", "RubricsFor") in _edge_labels(
result, "references", "attribute"
)
assert ("RubricsFor", "RubricFor") in _edge_labels(
result, "references", "return_type"
)
assert not [
node
for node in result["nodes"]
if node.get("label") in {"RubricFor", "RubricsFor"}
and not node.get("source_file")
]
def test_java_annotation_member_keeps_qualified_type_identity(tmp_path):
from graphify.extract import extract
internal = tmp_path / "internal" / "Rule.java"
internal.parent.mkdir()
internal.write_text("package internal;\npublic class Rule {}\n")
outer = tmp_path / "internal" / "Outer.java"
outer.write_text(
"package internal;\n"
"public class Outer { public static class Nested {} }\n"
)
local = tmp_path / "app" / "Rule.java"
local.parent.mkdir()
local.write_text("package app;\npublic class Rule {}\n")
annotation = tmp_path / "app" / "Uses.java"
annotation.write_text(
"package app;\n"
"public @interface Uses {\n"
" internal.Rule[] direct();\n"
" Class<internal.Rule> generic();\n"
" internal.Outer.Nested[] nestedDirect();\n"
" Class<internal.Outer.Nested> nestedGeneric();\n"
" java.lang.String label();\n"
"}\n"
)
result = extract(
[internal, outer, local, annotation], cache_root=tmp_path / "cache"
)
by_id = {node["id"]: node for node in result["nodes"]}
targets = [
(by_id[edge["target"]], edge.get("context"))
for edge in result["edges"]
if edge.get("relation") == "references"
and edge.get("source_file", "").endswith("Uses.java")
and by_id[edge["source"]].get("label") == "Uses"
and by_id[edge["target"]].get("label") == "Rule"
]
assert {
context
for target, context in targets
if target.get("source_file", "").endswith("internal/Rule.java")
} == {"return_type", "generic_arg"}
assert not [
target
for target, _context in targets
if target.get("source_file", "").endswith("app/Rule.java")
]
assert not [
node
for node in result["nodes"]
if node.get("label") == "java.lang.String"
]
assert {
context
for target, context in [
(by_id[edge["target"]], edge.get("context"))
for edge in result["edges"]
if edge.get("relation") == "references"
and edge.get("source_file", "").endswith("Uses.java")
and by_id[edge["source"]].get("label") == "Uses"
and by_id[edge["target"]].get("label") == "Nested"
]
if target.get("source_file", "").endswith("Outer.java")
} == {"return_type", "generic_arg"}
def test_java_enum_and_annotation_declarations_are_type_nodes(tmp_path):
source = tmp_path / "TypeDeclarations.java"
source.write_text(