1 files changed, 37 insertions, 14 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 8ff91a0..6bfd0be 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
@@ -8,6 +8,7 @@ import java.util.List;
 import java.util.Set;
 import java.util.TreeSet;
 
+import static java.lang.reflect.Modifier.*;
 import static org.springframework.util.StringUtils.collectionToDelimitedString;
 
 public class ClassG {
@@ -19,6 +20,10 @@ public class ClassG {
     private final List<MethodRef> methods = new ArrayList<>();
     private final List<Constructor> constructors = new ArrayList<>();
 
+    public ClassG(String packageName, String className) {
+        this(packageName, className, null);
+    }
+
     public ClassG(String packageName, String className, String extendsClass) {
         this.packageName = packageName;
         this.className = className;
@@ -64,20 +69,32 @@ public class ClassG {
         return ref;
     }
 
-    public FieldRef addField(Class<?> klass, String name) {
+    public FieldRef addField(TypeMirror klass, String name) {
         TypeRef type = addImport(klass);
-        FieldRef ref = new FieldRef(Modifier.PRIVATE | Modifier.FINAL, type, name);
+        FieldRef ref = new FieldRef(PRIVATE | FINAL, type, name);
         fields.add(ref);
         return ref;
     }
 
-    public FieldRef addPublicFinalField(Class<?> klass, String name) {
+    public FieldRef addField(int modifiers, Class<?> klass, String name) {
         TypeRef type = addImport(klass);
-        FieldRef ref = new FieldRef(Modifier.PUBLIC | Modifier.FINAL, type, name);
+        FieldRef ref = new FieldRef(modifiers, type, name);
         fields.add(ref);
         return ref;
     }
 
+    public FieldRef addField(Class<?> klass, String name) {
+        return addField(PRIVATE | FINAL, klass, name);
+    }
+
+    public FieldRef addPublicFinalField(Class<?> klass, String name) {
+        return addField(PUBLIC | FINAL, klass, name);
+    }
+
+    public FieldRef addPublicStaticFinalField(Class<?> klass, String name) {
+        return addField(PUBLIC | STATIC | FINAL, klass, name);
+    }
+
     public Constructor addConstructor(Parameters parameters, List<String> body) {
         Constructor constructor = new Constructor(this, parameters, body);
         constructors.add(constructor);
@@ -85,13 +102,19 @@ public class ClassG {
     }
 
     public MethodRef addMethod(List<String> body, TypeRef returnType, String name, ParameterRef... parameters) {
-        MethodRef ref = new MethodRef(returnType, name, parameters, body);
+        MethodRef ref = new MethodRef(PUBLIC, returnType, name, parameters, body);
+        methods.add(ref);
+        return ref;
+    }
+
+    public MethodRef addStaticMethod(List<String> body, TypeRef returnType, String name, ParameterRef... parameters) {
+        MethodRef ref = new MethodRef(PUBLIC | STATIC, returnType, name, parameters, body);
         methods.add(ref);
         return ref;
     }
 
     public void write(PrintWriter writer) {
-        if(packageName.length() > 0) {
+        if (packageName.length() > 0) {
             writer.println("package " + packageName + ";");
             writer.println();
         }
@@ -99,12 +122,14 @@ public class ClassG {
             if (i.isPrimitive()) {
                 continue;
             }
-            if(i.canonicalName().indexOf('.') == -1) {
+            if (i.canonicalName().indexOf('.') == -1) {
                 continue;
             }
             writer.println("import " + i.canonicalName() + ";");
         }
-        writer.println();
+        if (!imports.isEmpty()) {
+            writer.println();
+        }
 
         String extendsString = extendsClass == null ? "" : " extends " + extendsClass;
 
@@ -143,19 +168,17 @@ public class ClassG {
         }
 
         writer.print("    public " + returnString + " " + method.name + "(" + collectionToDelimitedString(parameters, ", ") + ")");
-        if(method.exceptions.isEmpty()) {
+        if (method.exceptions.isEmpty()) {
             writer.println(" {");
-        }
-        else {
+        } else {
             writer.println(" throws");
             ArrayList<TypeRef> typeRefs = new ArrayList<>(method.exceptions);
             for (int i = 0; i < typeRefs.size(); i++) {
                 TypeRef e = typeRefs.get(i);
                 writer.print("        " + e.name);
-                if(i < typeRefs.size() - 1) {
+                if (i < typeRefs.size() - 1) {
                     writer.println(",");
-                }
-                else {
+                } else {
                     writer.println(" {");
                 }
             }