Create custom annotation for Lombok

22.5k views Asked by At

I have used Lombok in my code to automatically generate getter and setter code. I want to add other personal annotations and use it.

For example, I want to add an @Exist method which verifies the existence of a key in a list:

@Getter    @Setter
public class User {

    private String name;
    private List<Integer> keys;

    public boolean existKeys(Integer key) {
        boolean exist = keys.contains(key);
        return exist;
    }
}

After creating the annotation, I would do something like:

@Getter    @Setter
public class User {

    private String name;
    @Exist
    private List<Integer> keys;
} 
1

There are 1 answers

10
Fedor Losev On BEST ANSWER

General Considerations

If you are already using Lombok, you can add custom Lombok transformation annotation and handler.

  1. Define Exists annotation with @Target(FIELD) and @Retention(SOURCE)

  2. Create a handler

    @ProviderFor(JavacAnnotationHandler.class)
    public class HandleExists extends JavacAnnotationHandler<Exists>{ ...` 
    

    to process your annotation. Handler class package must start with the lombok. prefix. If you need to support Eclipse, etc. in addition to javac, you'll need to write more handlers extending appropriate framework classes.

  3. In the handler override/implement the handle() method to generate the required code through AST manipulation.


You can take as a sample the @Getter implementation:

Annotation: Getter.java

Handler: HandleGetter.java

You can also look into sources of other annotations and handlers to see how to generate particular code.

You'll need to add dependencies on lombok, JDK tools.jar.


Some resources:


Note, there are some points to consider here

  • This is a bunch of non-trivial code to write and maintain. If you plan to use annotation 5-6 times it is just not worth it.
  • You may need to change your annotation processor implementation with lombok upgrades.
  • The hole in compiler that lombok relies on also may be closed (then the whole Lombok project will change dramatically or cease to exist; in this case you'll have a more serious problem anyway if you use Lombok extensively, even if just for @Getter).

A more complex alternative without Lombok is to use standard annotation processing for code generation but, AFAIK, you can't change original classes and must generate/use classes that extend them (unless you'll exploit the same back-door as Lombok or resort to a code manipulation like CGLib or ASM).


Lombok Example

Below is some working code to create custom Lombok annotation that I've called @Contains.

It is javac implementation only, no Eclipse, etc. I guess it will be not hard to create a similar handler for Eclipse or other IDE.

It will generate fieldNameContains() member method which is delegated to the fieldName.contains().

Note, the code is just quick and dirty (but working) sample. For production grade annotation, you will need to handle many boundary conditions, check correct types, handle Lombok configuration and so on, as it can be observed in lombok or lombok-pg library sources.


Sample usage


SomeEnity.java

@Getter
@Setter
public class SomeEntity {

    @NonNull
    @Contains
    private Collection<String> fieldOne = new ArrayList<>();

    @NonNull
    @Contains
    private Collection<String> fieldTwo = new ArrayList<>();

}

SomeEntityTest.java

public class SomeEntityTest {

    @Test
    public void test() {
        SomeEntity entity = new SomeEntity();

        Collection<String> test1 = Arrays.asList(new String[] { "1", "2" });
        entity.setFieldOne(test1);
        assertSame(test1, entity.getFieldOne());

        Collection<String> test2 = new HashSet<String>(Arrays.asList(new String[] { "3", "4" }));
        entity.setFieldTwo(test2);
        assertSame(test2, entity.getFieldTwo());

        assertTrue(entity.fieldOneContains("1"));
        assertTrue(entity.fieldOneContains("2"));
        assertFalse(entity.fieldOneContains("3"));
        assertFalse(entity.fieldOneContains("4"));

        assertFalse(entity.fieldTwoContains("1"));
        assertFalse(entity.fieldTwoContains("2"));
        assertTrue(entity.fieldTwoContains("3"));
        assertTrue(entity.fieldTwoContains("4"));

        try {
            entity.setFieldOne(null);
            fail("exception expected");
        } catch (Exception ex) {
        }

        try {
            entity.setFieldTwo(null);
            fail("exception expected");
        } catch (Exception ex) {
        }

    }
}

Annotation Implementaiton


Contains.java

@Target({ElementType.FIELD})
@Retention(RetentionPolicy.SOURCE)
public @interface Contains {
    Class<?>[] types() default {};
    Class<?>[] excludes() default {};
}

HandleContains.java

@ProviderFor(JavacAnnotationHandler.class) 
@HandlerPriority(65536) 
@ResolutionResetNeeded 
public class HandleContains extends JavacAnnotationHandler<Contains> {
    
    @Override 
    public void handle(AnnotationValues<Contains> annotation, JCAnnotation ast, JavacNode annotationNode) {
        
        try {
            JavacNode node = annotationNode.up();
            if (node.getKind() != Kind.FIELD) {
                annotationNode.addError("@Contains is allowed only on fields");
                return;
            }
            Name delegateName = annotationNode.toName(node.getName());
            JavacResolution reso = new JavacResolution(annotationNode.getContext());
            JCTree member = node.get();
            if (member.type == null) {
                reso.resolveClassMember(node);
            }
            Type delegateType = member.type;
            if (delegateType instanceof ClassType) {
                ClassType ct = (ClassType) delegateType;
                //TODO validate that this field is a collection type
                // if(!Collection)
                //   annotationNode.addError("@Contains can only be used on collections");
                final String methodName = "contains";
                MethodSig methodSig = getMethodBinding(methodName, ct, annotationNode.getTypesUtil());
                if (methodSig == null) throw new Exception("no method " + methodName + " in " + ct.tsym.name);
                JCMethodDecl methodDecl = createDelegateMethod(methodSig, annotationNode, delegateName);
                injectMethod(node.up(), methodDecl);
            } else {
                annotationNode.addError("@Contains can only use concrete class types");
                return;
            }
        } catch (Exception ex) {
            //ex.printStackTrace();
            annotationNode.addError("@Contains unexpected error: " + ex.getMessage());
        }
        
    }
    
    public JCMethodDecl createDelegateMethod(MethodSig sig, JavacNode annotation, Name delegateName) throws TypeNotConvertibleException {
        
        JavacTreeMaker maker = annotation.getTreeMaker();
        
        com.sun.tools.javac.util.List<JCAnnotation> annotations;
        if (sig.isDeprecated) {
            annotations = com.sun.tools.javac.util.List.of(maker.Annotation(genJavaLangTypeRef(annotation, "Deprecated"), com.sun.tools.javac.util.List.<JCExpression>nil()));
        } else {
            annotations = com.sun.tools.javac.util.List.nil();
        }
        
        JCModifiers mods = maker.Modifiers(PUBLIC, annotations);
        JCExpression returnType = JavacResolution.typeToJCTree((Type) sig.type.getReturnType(), annotation.getAst(), true);
        boolean useReturn = sig.type.getReturnType().getKind() != TypeKind.VOID;
        ListBuffer<JCVariableDecl> params = sig.type.getParameterTypes().isEmpty() ? null : new ListBuffer<JCVariableDecl>();
        ListBuffer<JCExpression> args = sig.type.getParameterTypes().isEmpty() ? null : new ListBuffer<JCExpression>();
        ListBuffer<JCExpression> thrown = sig.type.getThrownTypes().isEmpty() ? null : new ListBuffer<JCExpression>();
        ListBuffer<JCTypeParameter> typeParams = sig.type.getTypeVariables().isEmpty() ? null : new ListBuffer<JCTypeParameter>();
        ListBuffer<JCExpression> typeArgs = sig.type.getTypeVariables().isEmpty() ? null : new ListBuffer<JCExpression>();
        Types types = Types.instance(annotation.getContext());
        
        for (TypeMirror param : sig.type.getTypeVariables()) {
            Name name = ((TypeVar) param).tsym.name;
            
            ListBuffer<JCExpression> bounds = new ListBuffer<JCExpression>();
            for (Type type : types.getBounds((TypeVar) param)) {
                bounds.append(JavacResolution.typeToJCTree(type, annotation.getAst(), true));
            }
            
            typeParams.append(maker.TypeParameter(name, bounds.toList()));
            typeArgs.append(maker.Ident(name));
        }
        
        for (TypeMirror ex : sig.type.getThrownTypes()) {
            thrown.append(JavacResolution.typeToJCTree((Type) ex, annotation.getAst(), true));
        }
        
        int idx = 0;
        String[] paramNames = sig.getParameterNames();
        boolean varargs = sig.elem.isVarArgs();
        for (TypeMirror param : sig.type.getParameterTypes()) {
            long flags = JavacHandlerUtil.addFinalIfNeeded(Flags.PARAMETER, annotation.getContext());
            JCModifiers paramMods = maker.Modifiers(flags);
            Name name = annotation.toName(paramNames[idx++]);
            if (varargs && idx == paramNames.length) {
                paramMods.flags |= VARARGS;
            }
            params.append(maker.VarDef(paramMods, name, JavacResolution.typeToJCTree((Type) param, annotation.getAst(), true), null));
            args.append(maker.Ident(name));
        }
        
        JCExpression accessor = maker.Select(maker.Ident(annotation.toName("this")), delegateName);
        
        JCExpression delegateCall = maker.Apply(toList(typeArgs), maker.Select(accessor, sig.name), toList(args));
        JCStatement body = useReturn ? maker.Return(delegateCall) : maker.Exec(delegateCall);
        JCBlock bodyBlock = maker.Block(0, com.sun.tools.javac.util.List.of(body));
        StringBuilder generatedMethodName = new StringBuilder(delegateName);
        generatedMethodName.append(sig.name.toString());
        generatedMethodName.setCharAt(delegateName.length(), Character.toUpperCase(generatedMethodName.charAt(delegateName.length())));
        return recursiveSetGeneratedBy(maker.MethodDef(mods, annotation.toName(generatedMethodName.toString()), returnType, toList(typeParams), toList(params), toList(thrown), bodyBlock, null), annotation.get(), annotation.getContext());
    }
    
    public static <T> com.sun.tools.javac.util.List<T> toList(ListBuffer<T> collection) {
        return collection == null ? com.sun.tools.javac.util.List.<T>nil() : collection.toList();
    }
    
    public static class MethodSig {
        final Name name;
        final ExecutableType type;
        final boolean isDeprecated;
        final ExecutableElement elem;
        
        MethodSig(Name name, ExecutableType type, boolean isDeprecated, ExecutableElement elem) {
            this.name = name;
            this.type = type;
            this.isDeprecated = isDeprecated;
            this.elem = elem;
        }
        
        String[] getParameterNames() {
            List<? extends VariableElement> paramList = elem.getParameters();
            String[] paramNames = new String[paramList.size()];
            for (int i = 0; i < paramNames.length; i++) {
                paramNames[i] = paramList.get(i).getSimpleName().toString();
            }
            return paramNames;
        }
        
        @Override public String toString() {
            return (isDeprecated ? "@Deprecated " : "") + name + " " + type;
        }
    }
    
    public MethodSig getMethodBinding(String name, ClassType ct, JavacTypes types) {
        MethodSig result = null;
        TypeSymbol tsym = ct.asElement();
        if (tsym == null) throw new IllegalArgumentException("no class");
        
        for (Symbol member : tsym.getEnclosedElements()) {
            if (member.getKind() != ElementKind.METHOD || !name.equals(member.name.toString())) {
                continue;
            }
            if (member.isStatic()) continue;
            if (member.isConstructor()) continue;
            ExecutableElement exElem = (ExecutableElement) member;
            if (!exElem.getModifiers().contains(Modifier.PUBLIC)) continue;
            ExecutableType methodType = (ExecutableType) types.asMemberOf(ct, member);
            boolean isDeprecated = (member.flags() & DEPRECATED) != 0;
            result = new MethodSig(member.name, methodType, isDeprecated, exElem);
        }
        if (result == null) {
            if (ct.supertype_field instanceof ClassType) {
                result = getMethodBinding(name, (ClassType) ct.supertype_field, types);
            }
            if (result == null) {
                if (ct.interfaces_field != null) {
                    for (Type iface : ct.interfaces_field) {
                        if (iface instanceof ClassType) {
                            result = getMethodBinding(name, (ClassType) iface, types);
                            if (result != null) {
                                break;
                            }
                        }
                    }
                }
            }
        }
        return result;
    }
}