1 files changed, 10 insertions, 8 deletions

diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
index fbd7660..f944b1d 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
@@ -2,7 +2,6 @@ package io.trygvis.container.compiler;
 
 import io.trygvis.container.compiler.model.ClassG;
 import io.trygvis.container.compiler.model.FieldRef;
-import io.trygvis.container.compiler.model.ParameterRef;
 import io.trygvis.container.compiler.model.Parameters;
 import io.trygvis.container.compiler.model.TypeRef;
 import io.trygvis.container.tx.PlatformTransactionManager;
@@ -23,6 +22,7 @@ import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.List;
 
+import static java.lang.reflect.Modifier.PUBLIC;
 import static javax.lang.model.util.ElementFilter.constructorsIn;
 import static org.springframework.util.StringUtils.collectionToDelimitedString;
 
@@ -42,8 +42,8 @@ public class TransactionalHandler extends AbstractHandler {
 
         String p = elements.getPackageOf(element).getQualifiedName().toString();
 
-        ClassG g = new ClassG(new TypeRef(p + "." + className), new TypeRef(targetClassName.toString()));
-        FieldRef transactionManager = g.addField(PlatformTransactionManager.class, "transactionManager");
+        ClassG g = new ClassG(PUBLIC, new TypeRef(p + "." + className)).extendsType(new TypeRef(targetClassName.toString()));
+        FieldRef transactionManager = g.addField(new TypeRef(PlatformTransactionManager.class), "transactionManager");
 
         for (ExecutableElement constructor : constructorsIn(elements.getAllMembers(element))) {
             if (!constructor.getModifiers().contains(Modifier.PUBLIC)) {
@@ -79,12 +79,12 @@ public class TransactionalHandler extends AbstractHandler {
             boolean isVoid = kind == TypeKind.VOID;
             TypeRef returnType = isVoid ? TypeRef.VOID : g.addImport(returnTypeMirror);
 
-            List<ParameterRef> parameters = new ArrayList<>();
+            Parameters parameters = new Parameters();
             List<String> arguments = new ArrayList<>();
             for (VariableElement ve : ee.getParameters()) {
 //                    parameters.add("final " + ve.asType().toString() + " " + ve.getSimpleName().toString());
                 TypeRef k = g.addImport(ve.asType());
-                parameters.add(new ParameterRef(k, ve.getSimpleName().toString()));
+                parameters.addParameter(k, ve.getSimpleName().toString());
                 arguments.add(ve.getSimpleName().toString());
             }
 
@@ -107,12 +107,14 @@ public class TransactionalHandler extends AbstractHandler {
             body.add("            }");
             body.add("        });");
 
-            g.addMethod(body, returnType, ee.getSimpleName().toString(), parameters.toArray(new ParameterRef[arguments.size()]));
+            g.addMethod(body, returnType, ee.getSimpleName().toString(), parameters);
         }
 
         JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile((p.length() == 0 ? "" : p + ".") + className, element);
         try (PrintWriter w = new PrintWriter(sourceFile.openWriter())) {
-            g.write(w);
+            for (String s : g.generate()) {
+                w.println(s);
+            }
         }
     }
 
@@ -121,7 +123,7 @@ public class TransactionalHandler extends AbstractHandler {
         List<String> goesToSuper = new ArrayList<>();
 
         TypeRef k = g.addImport(PlatformTransactionManager.class);
-        ParameterRef transactionManager = parameters.addParameter(k, "transactionManager");
+        Parameters.ParameterRef transactionManager = parameters.addParameter(k, "transactionManager");
 
         for (VariableElement p : constructor.getParameters()) {
             k = g.addImport(p.asType());