diff --git a/src/main/java/com/alibaba/testable/processor/TestableProcessor.java b/src/main/java/com/alibaba/testable/processor/TestableProcessor.java index 0f37390..165550e 100644 --- a/src/main/java/com/alibaba/testable/processor/TestableProcessor.java +++ b/src/main/java/com/alibaba/testable/processor/TestableProcessor.java @@ -1,9 +1,11 @@ package com.alibaba.testable.processor; import com.alibaba.testable.annotation.Testable; -import com.squareup.javapoet.JavaFile; -import com.squareup.javapoet.MethodSpec; -import com.squareup.javapoet.TypeSpec; +import com.alibaba.testable.translator.TestableTreeTranslator; +import com.alibaba.testable.util.ConstPool; +import com.squareup.javapoet.*; +import com.sun.tools.javac.code.Type; +import com.sun.tools.javac.tree.JCTree; import javax.annotation.processing.RoundEnvironment; import javax.annotation.processing.SupportedAnnotationTypes; @@ -15,6 +17,8 @@ import javax.lang.model.element.TypeElement; import javax.tools.JavaFileObject; import java.io.IOException; import java.io.Writer; +import java.util.ArrayList; +import java.util.List; import java.util.Set; /** @@ -45,28 +49,48 @@ public class TestableProcessor extends BaseProcessor { } private String createTestableClass(Element classElement, String packageName, String className) { - MethodSpec main = MethodSpec.methodBuilder("main") - .addModifiers(Modifier.PUBLIC, Modifier.STATIC) - .returns(void.class) - .addParameter(String[].class, "args") - .addStatement("$T.out.println($S)", System.class, "Hello, Testable !") - .build(); - MethodSpec constructor = MethodSpec.constructorBuilder() - .addModifiers(Modifier.PUBLIC) - .build(); + JCTree tree = trees.getTree(classElement); + TestableTreeTranslator translator = new TestableTreeTranslator(); + tree.accept(translator); - TypeSpec testableClass = TypeSpec.classBuilder(className) + List methodSpecs = new ArrayList<>(); + for (JCTree.JCMethodDecl method : translator.getMethods()) { + if (method.name.toString().equals(ConstPool.CONSTRUCTOR_NAME)) { + MethodSpec.Builder builder = MethodSpec.constructorBuilder() + .addModifiers(Modifier.PUBLIC); + for (JCTree.JCVariableDecl p : method.getParameters()) { + builder.addParameter(getParameterSpec(p.type)); + } + methodSpecs.add(builder.build()); + } else { + MethodSpec.Builder builder = MethodSpec.methodBuilder(method.name.toString()) + .addModifiers(method.getModifiers().getFlags()) + .addModifiers(Modifier.PUBLIC) + .returns(method.restype.getClass()); + for (JCTree.JCVariableDecl p : method.getParameters()) { + builder.addParameter(getParameterSpec(p.type)); + } + builder.addStatement("$T.out.println($S)", System.class, "Hello, Testable !"); + methodSpecs.add(builder.build()); + } + } + + TypeSpec.Builder builder = TypeSpec.classBuilder(className) .addModifiers(Modifier.PUBLIC, Modifier.FINAL) - .superclass(classElement.asType()) - .addMethod(constructor) - .addMethod(main) - .build(); - - JavaFile javaFile = JavaFile.builder(packageName, testableClass) - .build(); - + .superclass(classElement.asType()); + for (MethodSpec m : methodSpecs) { + builder.addMethod(m); + } + TypeSpec testableClass = builder.build(); + JavaFile javaFile = JavaFile.builder(packageName, testableClass).build(); return javaFile.toString(); } + private ParameterSpec getParameterSpec(Type type) { + return ParameterSpec.builder(String.class, "placeholder") + .addModifiers(Modifier.PUBLIC) + .build(); + } + } diff --git a/src/main/java/com/alibaba/testable/translator/TestableTreeTranslator.java b/src/main/java/com/alibaba/testable/translator/TestableTreeTranslator.java new file mode 100644 index 0000000..08ef961 --- /dev/null +++ b/src/main/java/com/alibaba/testable/translator/TestableTreeTranslator.java @@ -0,0 +1,38 @@ +package com.alibaba.testable.translator; + +import com.sun.tools.javac.tree.JCTree; +import com.sun.tools.javac.tree.JCTree.JCMethodDecl; +import com.sun.tools.javac.tree.TreeTranslator; +import com.sun.tools.javac.util.List; + +import java.lang.reflect.Modifier; + +/** + * Travel AST + * + * @author flin + */ +public class TestableTreeTranslator extends TreeTranslator { + + /** + * Methods to inject + */ + private List methods = List.nil(); + + public List getMethods() { + return methods; + } + + @Override + public void visitClassDef(JCTree.JCClassDecl jcClassDecl) { + super.visitClassDef(jcClassDecl); + jcClassDecl.mods.flags = jcClassDecl.mods.flags & (~Modifier.FINAL); + } + + @Override + public void visitMethodDef(JCMethodDecl jcMethodDecl) { + super.visitMethodDef(jcMethodDecl); + methods = methods.append(jcMethodDecl); + } + +} diff --git a/src/main/java/com/alibaba/testable/util/ConstPool.java b/src/main/java/com/alibaba/testable/util/ConstPool.java index cd12f19..111b279 100644 --- a/src/main/java/com/alibaba/testable/util/ConstPool.java +++ b/src/main/java/com/alibaba/testable/util/ConstPool.java @@ -5,4 +5,6 @@ package com.alibaba.testable.util; */ public final class ConstPool { + public static final String CONSTRUCTOR_NAME = ""; + }