【Java进阶】利用APT优雅的实现统一日志格式

/**
 * 用于自动生成log成员变量.仅适用于class或enum,不适用于接口.
 */
@Retention(RetentionPolicy.SOURCE)
@Target(ElementType.TYPE)
public @interface Slf4j {

    /**
     * 系统名称.如果为空则取"-Dvlogging.system"系统属性,如果系统属性也为空,则取"Unknown".
     */
    String system() default "";

    /**
     * 模块名称.如果为空则取"-Dvlogging.module"系统属性,如果系统属性也为空,则取"Unknown".
     */
    String module() default "";
}

public class Slf4jProcessor extends AbstractProcessor {

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        messager = processingEnv.getMessager();
        trees = Trees.instance(processingEnv);
        Context context = ((JavacProcessingEnvironment) processingEnv).getContext();
        maker = TreeMaker.instance(context);
        names = Names.instance(context);
    }

    ...
}

@Override
public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
    // 1 检查类型
    roundEnv.getElementsAnnotatedWith(Slf4j.class).stream().forEach(elm -> {
        if (elm.getKind() != ElementKind.CLASS && elm.getKind() != ElementKind.ENUM) {
            messager.printMessage(Diagnostic.Kind.ERROR, "Only classes or enums can be annotated with " + Slf4j.class.getSimpleName());
            return;
        }

        // 2 检查log成员变量是否已存在
        TypeElement typeElm = (TypeElement) elm;
        if (typeElm.getEnclosedElements().stream()
                .filter(e -> e.getKind() == ElementKind.FIELD && Logger.FIELD_NAME.equals(e.getSimpleName())).count() > 0) {
            messager.printMessage(Diagnostic.Kind.WARNING, MessageFormat.format("A member field named {0} already exists in the annotated class", Logger.FIELD_NAME));
            return;
        }

        // 3 注入log成员变量
        CompilationUnitTree cuTree = trees.getPath(typeElm).getCompilationUnit();
        if (cuTree instanceof JCTree.JCCompilationUnit) {
            JCTree.JCCompilationUnit cu = (JCTree.JCCompilationUnit) cuTree;
            // only process on files which have been compiled from source
            if (cu.sourcefile.getKind() == JavaFileObject.Kind.SOURCE) {
                _findType(cu, typeElm.getQualifiedName().toString()).ifPresent(type -> {
                    Slf4j slf4j = typeElm.getAnnotation(Slf4j.class);
                    String system = slf4j.system();
                    String module = slf4j.module();

                    // 生成private static final Logger log = LoggerFactory.getLogger(LoggerFactory.Type.SLF4J, <annotatedClass>, <system>, <module>);
                    JCTree.JCExpression loggerType = _toExpression(Logger.class.getCanonicalName());
                    JCTree.JCExpression getLoggerMethod = _toExpression(LoggerFactory.class.getCanonicalName() + ".getLogger");
                    JCTree.JCExpression typeArg = _toExpression(LoggerFactory.Type.class.getCanonicalName() + "." + LoggerFactory.Type.SLF4J.name());
                    JCTree.JCExpression nameArg = _toExpression(typeElm.getQualifiedName() + ".class");
                    JCTree.JCExpression systemArg = maker.Literal(system);
                    JCTree.JCExpression moduleArg = maker.Literal(module);
                    JCTree.JCMethodInvocation getLoggerCall = maker.Apply(List.nil(), getLoggerMethod, List.of(typeArg, nameArg, systemArg, moduleArg));
                    JCTree.JCVariableDecl logField = maker.VarDef(
                            maker.Modifiers(Flags.PRIVATE | Flags.STATIC | Flags.FINAL),
                            names.fromString(Logger.FIELD_NAME), loggerType, getLoggerCall);

                    _insertField(type, logField);
                });
            }
        }
    });

    return true;
}

@Slf4j(system = "Vlogging", module = "Integration")
public class VloggingAnnotated {

    public static void main(String[] args) {
        HashMap<String, String> params = new HashMap<>();
        params.put("foo", "xyz");
        log.info(VloggingAnnotated.class.getCanonicalName(), params);
    }
}

// 继承com.intellij.psi.augment.PsiAugmentProvider类

@NotNull
@Override
public <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement psiElement, @NotNull Class<Psi> type) {
    final List<Psi> emptyResult = Collections.emptyList();
    // skip processing during index rebuild
    final Project project = psiElement.getProject();
    if (DumbService.isDumb(project)) {
        return emptyResult;
    }
    // Expecting that we are only augmenting an PsiClass
    // Don't filter !isPhysical elements or code auto completion will not work
    if (!(psiElement instanceof PsiExtensibleClass) || !psiElement.isValid()) {
        return emptyResult;
    }
    // filter non-field type
    if (!PsiField.class.isAssignableFrom(type)) {
        return emptyResult;
    }
    final PsiClass psiClass = (PsiClass) psiElement;
    // see AbstractClassProcessor#process()
    PsiAnnotation psiAnnotation = PsiAnnotationUtil.findAnnotation(psiClass, Slf4j.class);
    if (null == psiAnnotation) {
        return emptyResult;
    }
    // check cache first
    if (loggerCache.containsKey(psiClass.getQualifiedName())) {
        return Arrays.asList((Psi) loggerCache.get(psiClass.getQualifiedName()));
    }

    final PsiManager manager = psiClass.getContainingFile().getManager();
    final PsiElementFactory psiElementFactory = JavaPsiFacade.getElementFactory(project);
    PsiType psiLoggerType = psiElementFactory.createTypeFromText(LOGGER_TYPE, psiClass);
    LightFieldBuilder loggerField = new LightFieldBuilder(manager, LOGGER_NAME, psiLoggerType);
    LightModifierList modifierList = (LightModifierList) loggerField.getModifierList();
    modifierList.addModifier(PsiModifier.PRIVATE);
    modifierList.addModifier(PsiModifier.STATIC);
    modifierList.addModifier(PsiModifier.FINAL);
    loggerField.setContainingClass(psiClass);
    loggerField.setNavigationElement(psiAnnotation);

    final String loggerInitializerParameter = String.format(LOGGER_CATEGORY, psiClass.getName());
    final PsiExpression initializer = psiElementFactory.createExpressionFromText(String.format(LOGGER_INITIALIZER, loggerInitializerParameter), psiClass);
    loggerField.setInitializer(initializer);
    // add to cache
    loggerCache.put(psiClass.getQualifiedName(), loggerField);

    return Arrays.asList((Psi) loggerField);
}