1 files changed, 25 insertions, 9 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
index c99feae..402498c 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
@@ -45,7 +45,12 @@ public class ClassG {
     }
 
     public ClassG extendsType(TypeRef extendsType) {
-        this.extendsType = extendsType;
+        this.extendsType = imports.add(extendsType);
+        return this;
+    }
+
+    public ClassG extendsType(Class<?> extendsType) {
+        this.extendsType = imports.add(extendsType);
         return this;
     }
 
@@ -54,6 +59,13 @@ public class ClassG {
         return this;
     }
 
+    public ClassG implementsType(Class... implementsTypes) {
+        for (Class type : implementsTypes) {
+            this.implementsTypes.add(imports.add(type));
+        }
+        return this;
+    }
+
     public FieldRef addField(TypeMirror klass, String name) {
         TypeRef type = imports.add(klass);
         FieldRef ref = new FieldRef(PRIVATE | FINAL, type, name);
@@ -80,18 +92,17 @@ public class ClassG {
         return addField(PUBLIC | STATIC | FINAL, type, name);
     }
 
-    public Constructor addConstructor(Parameters parameters, List<String> body) {
-        Constructor constructor = new Constructor(this, parameters, body);
+    public ClassG add(Constructor constructor) {
         constructors.add(constructor);
-        return constructor;
+        return this;
     }
 
-    public ClassG addMethod(MethodRef methodRef) {
+    public ClassG add(MethodRef methodRef) {
         this.methods.add(methodRef);
         return this;
     }
 
-//    public MethodRef addMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
+//    public MethodRef add(List<String> body, TypeRef returnType, String name, Parameters parameters) {
 //        MethodRef ref = new MethodRef(PUBLIC, returnType, name, parameters, body);
 //        methods.add(ref);
 //        return ref;
@@ -145,7 +156,7 @@ public class ClassG {
 
         for (Constructor constructor : constructors) {
             body.add("");
-            addAll(1, body, constructor.write());
+            addAll(1, body, constructor.write(this));
         }
 
         for (MethodRef method : methods) {
@@ -165,7 +176,11 @@ public class ClassG {
     }
 
     private List<String> write(MethodRef method) {
-        List<String> body = new ArrayList<>();
+        String typeArgs = "";
+        if (!method.typeArgs.isEmpty()) {
+            typeArgs = "<" + join(method.typeArgs, ", ") + "> ";
+        }
+
         String returnString;
         if (method.returnType == TypeRef.VOID) {
             returnString = "void";
@@ -178,9 +193,10 @@ public class ClassG {
             parameters.add("final " + p.klass + " " + p.name);
         }
 
-        String m = Modifier.toString(method.modifiers) + " " +
+        String m = Modifier.toString(method.modifiers) + " " + typeArgs +
                 returnString + " " +
                 method.name + "(" + collectionToDelimitedString(parameters, ", ") + ")";
+        List<String> body = new ArrayList<>();
         if (method.exceptions.isEmpty()) {
             body.add(m + " {");
         } else {