Nits.

4 files changed, 29 insertions, 6 deletions

diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/AnnotationG.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/AnnotationG.java
new file mode 100644
index 0000000..4b6c05c
--- /dev/null
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/AnnotationG.java
@@ -0,0 +1,15 @@
+package io.trygvis.container.compiler.model;
+
+public class AnnotationG {
+    public final TypeRef type;
+    public final String body;
+
+    public AnnotationG(TypeRef type, String body) {
+        this.type = type;
+        this.body = body;
+    }
+
+    public String toJava() {
+        return "@" + type.plainName + "(" + body + ")";
+    }
+}
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 3aa9c62..fe595a0 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
@@ -9,7 +9,6 @@ import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
-import java.util.TreeSet;
 
 import static io.trygvis.container.compiler.Utils.toJavaString;
 import static java.lang.reflect.Modifier.*;
@@ -195,10 +194,13 @@ public class ClassG {
             parameters.add("final " + p.klass + " " + p.name);
         }
 
+        List<String> body = new ArrayList<>();
+        for (AnnotationG annotation : method.annotations) {
+            body.add(annotation.toJava());
+        }
         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 {
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/MethodRef.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/MethodRef.java
index 952a1e5..65a5ff2 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/MethodRef.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/MethodRef.java
@@ -17,6 +17,7 @@ public class MethodRef {
     public final Parameters parameters;
     public final Set<TypeRef> exceptions = new TreeSet<>();
     public final List<String> typeArgs = new ArrayList<>();
+    public final List<AnnotationG> annotations = new ArrayList<>();
     public final List<String> body;
 
     public MethodRef(int modifiers, TypeRef returnType, String name, String body) {
@@ -48,4 +49,9 @@ public class MethodRef {
         addAll(typeArgs, args);
         return this;
     }
+
+    public MethodRef annotation(AnnotationG annotation) {
+        this.annotations.add(annotation);
+        return this;
+    }
 }
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java
index f524b1a..f621343 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java
@@ -1,6 +1,7 @@
 package io.trygvis.persistence.generators;
 
 import io.trygvis.container.compiler.SqlUnitModel;
+import io.trygvis.container.compiler.model.AnnotationG;
 import io.trygvis.container.compiler.model.ClassG;
 import io.trygvis.container.compiler.model.Constructor;
 import io.trygvis.container.compiler.model.FieldRef;
@@ -66,13 +67,12 @@ public class EntityManagerGenerator {
         List<String> body = new ArrayList<>();
         for (EntityMirror entity : unit.getEntities().values()) {
             body.add("if (klass == " + entity.type.plainName + ".class) {");
-            body.add("    @SuppressWarnings(\"unchecked\")");
-            body.add("    SqlDao<Id, T> dao = (SqlDao<Id, T>) " + daoFields.get(entity).name + ";");
-            body.add("    return (SqlDao<Id, T>) dao;");
+            body.add("    return (SqlDao<Id, T>) " + daoFields.get(entity).name + ";");
             body.add("}");
         }
         body.add("throw new RuntimeException(\"Type is not a part of this persistence unit: \" + klass);");
-        return new MethodRef(PUBLIC, sqlDatoType, "getDao", p, body).typeArgs("Id", "T");
+        return new MethodRef(PUBLIC, sqlDatoType, "getDao", p, body).typeArgs("Id", "T").
+                annotation(new AnnotationG(new TypeRef(SuppressWarnings.class), "\"unchecked\""));
     }
 
     /*