2020import org .jspecify .annotations .Nullable ;
2121import org .openrewrite .*;
2222import org .openrewrite .java .*;
23- import org .openrewrite .java .service .AnnotationService ;
2423import org .openrewrite .java .tree .Expression ;
2524import org .openrewrite .java .tree .J ;
2625import org .openrewrite .java .tree .JavaType ;
2726import org .openrewrite .java .tree .TypeTree ;
2827import org .openrewrite .staticanalysis .java .MoveFieldAnnotationToType ;
2928
30- import java .util .Arrays ;
31- import java .util .Comparator ;
32- import java .util .List ;
33- import java .util .Optional ;
29+ import java .util .*;
30+ import java .util .Locale ;
3431import java .util .concurrent .atomic .AtomicBoolean ;
3532
3633@ EqualsAndHashCode (callSuper = false )
3734@ Value
3835public class AnnotateNullableMethods extends Recipe {
3936
37+ private static final String DEFAULT_NULLABLE_ANN_CLASS = "org.jspecify.annotations.Nullable" ;
38+
39+ /**
40+ * FQNs of nullable annotations that are meta-annotated with {@code @Target(TYPE_USE)}.
41+ * These annotations can be positioned before the inner type of a nested type or on array brackets.
42+ * All other nullable annotations are assumed to be declaration-target only and will remain
43+ * as method-level annotations.
44+ */
45+ private static final Set <String > TYPE_USE_NULLABLE_ANNOTATIONS = new HashSet <>(Arrays .asList (
46+ "jakarta.annotation.Nullable" ,
47+ "org.checkerframework.checker.nullness.qual.Nullable" ,
48+ "org.eclipse.jdt.annotation.Nullable" ,
49+ "org.jspecify.annotations.Nullable"
50+ ));
51+
4052 @ Option (displayName = "`@Nullable` annotation class" ,
41- description = "The fully qualified name of the @Nullable annotation. The annotation should be meta annotated with `@Target(TYPE_USE)`. Defaults to `org.jspecify.annotations.Nullable`" ,
53+ description = "The fully qualified name of the @Nullable annotation to add. " +
54+ "Both `@Target(TYPE_USE)` and declaration annotations (e.g. `javax.annotation.CheckForNull`) are supported. " +
55+ "Defaults to `org.jspecify.annotations.Nullable`." ,
4256 example = "org.jspecify.annotations.Nullable" ,
4357 required = false )
4458 @ Nullable
4559 String nullableAnnotationClass ;
4660
47- private static final String DEFAULT_NULLABLE_ANN_CLASS = "org.jspecify.annotations.Nullable" ;
48-
4961 String displayName = "Annotate methods which may return `null` with `@Nullable`" ;
5062
5163 String description = "Add `@Nullable` to non-private methods that may return `null`. " +
52- "By default `org.jspecify.annotations.Nullable` is used, but through the `nullableAnnotationClass` option a custom annotation can be provided. " +
53- "When providing a custom `nullableAnnotationClass` that annotation should be meta annotated with `@Target(TYPE_USE)`. " +
54- "This recipe scans for methods that do not already have a `@Nullable` annotation and checks their return " +
55- "statements for potential null values. It also identifies known methods from standard libraries that may " +
56- "return null, such as methods from `Map`, `Queue`, `Deque`, `NavigableSet`, and `Spliterator`. " +
57- "The return of streams, or lambdas are not taken into account." ;
64+ "By default `org.jspecify.annotations.Nullable` is used, but through the `nullableAnnotationClass` option a custom annotation can be provided. " +
65+ "Both `@Target(TYPE_USE)` and declaration annotations (e.g. `javax.annotation.CheckForNull`) are supported. " +
66+ "Methods that already carry a known nullable annotation (matched by simple name) are skipped to avoid duplication. " +
67+ "This recipe scans for methods that do not already have a `@Nullable` annotation and checks their return " +
68+ "statements for potential null values. It also identifies known methods from standard libraries that may " +
69+ "return null, such as methods from `Map`, `Queue`, `Deque`, `NavigableSet`, and `Spliterator`. " +
70+ "The return of streams, or lambdas are not taken into account." ;
5871
5972 @ Override
6073 public Validated <Object > validate () {
@@ -68,14 +81,15 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
6881 String fullyQualifiedName = nullableAnnotationClass != null ? nullableAnnotationClass : DEFAULT_NULLABLE_ANN_CLASS ;
6982 String fullyQualifiedPackage = fullyQualifiedName .substring (0 , fullyQualifiedName .lastIndexOf ('.' ));
7083 String simpleName = fullyQualifiedName .substring (fullyQualifiedName .lastIndexOf ('.' ) + 1 );
84+ boolean isTypeUseAnnotation = TYPE_USE_NULLABLE_ANNOTATIONS .contains (fullyQualifiedName );
85+
7186 JavaIsoVisitor <ExecutionContext > javaIsoVisitor = new JavaIsoVisitor <ExecutionContext >() {
7287 @ Override
7388 public J .MethodDeclaration visitMethodDeclaration (J .MethodDeclaration methodDeclaration , ExecutionContext ctx ) {
7489 if (!methodDeclaration .hasModifier (J .Modifier .Type .Public ) ||
7590 methodDeclaration .getMethodType () == null ||
7691 methodDeclaration .getMethodType ().getReturnType () instanceof JavaType .Primitive ||
77- service (AnnotationService .class ).matches (getCursor (), new AnnotationMatcher ("@" + fullyQualifiedName )) ||
78- hasNullableAnnotation (methodDeclaration .getReturnTypeExpression (), fullyQualifiedName )) {
92+ hasAnyNullableAnnotation (methodDeclaration )) {
7993 return methodDeclaration ;
8094 }
8195
@@ -88,33 +102,52 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration methodDecl
88102 .build ()
89103 .apply (getCursor (), md .getCoordinates ().addAnnotation (Comparator .comparing (J .Annotation ::getSimpleName )));
90104 doAfterVisit (ShortenFullyQualifiedTypeReferences .modifyOnly (annotatedMethod ));
91- doAfterVisit (new MoveFieldAnnotationToType (fullyQualifiedName ).getVisitor ());
92- return (J .MethodDeclaration ) new NullableOnMethodReturnType ().getVisitor ()
93- .visitNonNull (annotatedMethod , ctx , getCursor ().getParentTreeCursor ());
105+
106+ // TYPE_USE annotations are moved to the return type position (e.g. public @Nullable String foo())
107+ // and positioned before inner types of nested types (e.g. Outer.@Nullable Inner).
108+ // Declaration-target annotations stay as method-level annotations (e.g. @CheckForNull \n public String foo()).
109+ if (isTypeUseAnnotation ) {
110+ doAfterVisit (new MoveFieldAnnotationToType (fullyQualifiedName ).getVisitor ());
111+ return (J .MethodDeclaration ) new NullableOnMethodReturnType ().getVisitor ()
112+ .visitNonNull (annotatedMethod , ctx , getCursor ().getParentTreeCursor ());
113+ }
114+ return annotatedMethod ;
94115 }
95116 return md ;
96117 }
97118
98- private boolean hasNullableAnnotation (@ Nullable TypeTree returnType , String annotationFqn ) {
99- if (returnType == null ) {
100- return false ;
101- }
102-
103- // Check if the return type itself is annotated
104- if (service (AnnotationService .class ).matches (new Cursor (null , returnType ), new AnnotationMatcher ("@" + annotationFqn ))) {
105- return true ;
106- }
107-
108- // For array types, check if the element type is annotated
109- if (returnType instanceof J .ArrayType ) {
110- J .ArrayType arrayType = (J .ArrayType ) returnType ;
111- if (arrayType .getElementType () instanceof J .AnnotatedType ) {
112- return service (AnnotationService .class ).matches (new Cursor (null , arrayType .getElementType ()), new AnnotationMatcher ("@" + annotationFqn ));
119+ /**
120+ * Checks whether the method declaration already has any known nullable annotation,
121+ * either as a method-level annotation or anywhere on the return type.
122+ */
123+ private boolean hasAnyNullableAnnotation (J .MethodDeclaration methodDeclaration ) {
124+ // Check method-level annotations
125+ for (J .Annotation annotation : methodDeclaration .getLeadingAnnotations ()) {
126+ if (isNullAnnotation (annotation )) {
127+ return true ;
113128 }
114129 }
115-
130+ // Scan the entire return type tree for any annotation with a known nullable simple name.
131+ // Uses a TreeVisitor to reliably traverse all AST node types regardless of structure
132+ // (J.AnnotatedType, J.FieldAccess with annotated names, J.ArrayType with bracket annotations, etc.)
133+ TypeTree returnType = methodDeclaration .getReturnTypeExpression ();
134+ if (returnType != null ) {
135+ return new JavaIsoVisitor <AtomicBoolean >() {
136+ @ Override
137+ public J .Annotation visitAnnotation (J .Annotation annotation , AtomicBoolean found ) {
138+ if (isNullAnnotation (annotation )) {
139+ found .set (true );
140+ }
141+ return annotation ;
142+ }
143+ }.reduce (returnType , new AtomicBoolean (false ), getCursor ()).get ();
144+ }
116145 return false ;
117146 }
147+
148+ private boolean isNullAnnotation (J .Annotation ann ) {
149+ return ann .getSimpleName ().toLowerCase (Locale .ROOT ).contains ("null" );
150+ }
118151 };
119152 return Repeat .repeatUntilStable (javaIsoVisitor , 5 );
120153 }
0 commit comments