ROASTER-36: Fixed PropertyImpl.setType

Author: gastaldi

api/src/main/java/org/jboss/forge/roaster/model/source/PropertyHolderSource.java

@@ -9,6 +9,7 @@
 
 import java.util.List;
 
+import org.jboss.forge.roaster.model.JavaType;
 import org.jboss.forge.roaster.model.Property;
 import org.jboss.forge.roaster.model.PropertyHolder;
 
@@ -25,6 +26,16 @@ public interface PropertyHolderSource<O extends JavaSource<O>> extends
     */
    PropertySource<O> addProperty(String type, String name);
 
+   /**
+    * Add a new {@link Property} declaration to this {@link O} instance.
+    */
+   PropertySource<O> addProperty(Class<?> type, String name);
+
+   /**
+    * Add a new {@link Property} declaration to this {@link O} instance.
+    */
+   PropertySource<O> addProperty(JavaType<?> type, String name);
+
    /**
     * Remove the given {@link Property} from this {@link O} instance, if it exists; otherwise, do nothing.
     */

impl/src/main/java/org/jboss/forge/roaster/model/impl/AbstractJavaSourceMemberHolder.java

@@ -310,7 +310,8 @@ public boolean hasMethodSignature(final Method<?, ?> method)
             Iterator<ParameterSource<O>> localParams = local.getParameters().iterator();
             for (Parameter<? extends JavaType<?>> methodParam : method.getParameters())
             {
-               if (localParams.hasNext() && Strings.areEqual(localParams.next().getType().getName(), methodParam.getType().getName()))
+               if (localParams.hasNext()
+                        && Strings.areEqual(localParams.next().getType().getName(), methodParam.getType().getName()))
                {
                   continue;
                }
@@ -515,6 +516,18 @@ public org.jboss.forge.roaster.model.Type<O> getType()
       return getProperty(name);
    }
 
+   @Override
+   public PropertySource<O> addProperty(Class<?> type, String name)
+   {
+      return addProperty(type.getName(), name);
+   }
+
+   @Override
+   public PropertySource<O> addProperty(JavaType<?> type, String name)
+   {
+      return addProperty(type.getQualifiedName(), name);
+   }
+
    @Override
    public final AbstractJavaSourceMemberHolder<O> removeProperty(Property<O> property)
    {

impl/src/main/java/org/jboss/forge/roaster/model/impl/MethodImpl.java

@@ -16,7 +16,10 @@
 import org.eclipse.jdt.core.dom.Block;
 import org.eclipse.jdt.core.dom.CompilationUnit;
 import org.eclipse.jdt.core.dom.MethodDeclaration;
+import org.eclipse.jdt.core.dom.Name;
+import org.eclipse.jdt.core.dom.PrimitiveType;
 import org.eclipse.jdt.core.dom.Modifier.ModifierKeyword;
+import org.eclipse.jdt.core.dom.PrimitiveType.Code;
 import org.eclipse.jdt.core.dom.SimpleName;
 import org.eclipse.jdt.core.dom.SingleVariableDeclaration;
 import org.eclipse.jdt.core.dom.Statement;
@@ -248,13 +251,51 @@ public MethodSource<O> setReturnTypeVoid()
    @Override
    public MethodSource<O> setReturnType(final String typeName)
    {
-      String stub = "public class Stub { public " + typeName + " method() {} }";
-      JavaClassSource temp = (JavaClassSource) Roaster.parse(stub);
-      List<MethodSource<JavaClassSource>> methods = temp.getMethods();
-      org.eclipse.jdt.core.dom.Type returnType = ((MethodDeclaration) methods.get(0).getInternal()).getReturnType2();
+      String simpleName = Types.toSimpleName(typeName);
+
+      O origin = getOrigin();
+      if (!Strings.areEqual(typeName, simpleName) && origin.requiresImport(typeName))
+      {
+         origin.addImport(typeName);
+      }
 
-      returnType = (org.eclipse.jdt.core.dom.Type) ASTNode.copySubtree(method.getAST(), returnType);
-      method.setReturnType2(returnType);
+      Code primitive = PrimitiveType.toCode(typeName);
+
+      org.eclipse.jdt.core.dom.Type type = null;
+      if (primitive != null)
+      {
+         type = ast.newPrimitiveType(primitive);
+      }
+      else
+      {
+         if (!origin.requiresImport(typeName))
+         {
+            if (Types.isArray(typeName))
+            {
+               String arrayType = Types.stripArray(typeName);
+               int arrayDimension = Types.getArrayDimension(typeName);
+               if (Types.isPrimitive(arrayType))
+               {
+                  type = ast.newArrayType(ast.newPrimitiveType(PrimitiveType.toCode(arrayType)), arrayDimension);
+               }
+               else
+               {
+                  type = ast.newArrayType(ast.newSimpleType(ast.newSimpleName(arrayType)), arrayDimension);
+               }
+            }
+            else
+            {
+               type = ast.newSimpleType(ast.newSimpleName(simpleName));
+            }
+         }
+         else
+         {
+            String[] className = Types.tokenizeClassName(typeName);
+            Name name = ast.newName(className);
+            type = ast.newSimpleType(name);
+         }
+      }
+      method.setReturnType2(type);
 
       return this;
    }

impl/src/main/java/org/jboss/forge/roaster/model/impl/PropertyImpl.java

@@ -25,6 +25,7 @@
 import org.jboss.forge.roaster.model.source.FieldSource;
 import org.jboss.forge.roaster.model.source.JavaSource;
 import org.jboss.forge.roaster.model.source.MethodSource;
+import org.jboss.forge.roaster.model.source.ParameterSource;
 import org.jboss.forge.roaster.model.source.PropertyHolderSource;
 import org.jboss.forge.roaster.model.source.PropertySource;
 import org.jboss.forge.roaster.model.util.Assert;
@@ -339,7 +340,9 @@ public PropertySource<O> setType(String type)
       }
       if (mutator != null)
       {
-         mutator.setParameters(String.format("%s %s", type, getName()));
+         for (ParameterSource<O> param : mutator.getParameters())
+            mutator.removeParameter(param);
+         mutator.addParameter(type, getName());
       }
       if (field != null)
       {
@@ -450,7 +453,7 @@ public boolean equals(Object obj)
    @Override
    public int hashCode()
    {
-      // compatible with Java 6: 
+      // compatible with Java 6:
       return Arrays.hashCode(new Object[] { getOrigin(), getName() });
    }
 

impl/src/test/java/org/jboss/forge/test/roaster/model/JavaClassPropertyTest.java

@@ -8,6 +8,8 @@
 package org.jboss.forge.test.roaster.model;
 
 import java.io.InputStream;
+import java.sql.Timestamp;
+import java.util.Date;
 import java.util.List;
 
 import org.jboss.forge.roaster.Roaster;
@@ -59,4 +61,16 @@ public void testBooleanProperties()
       Assert.assertNotNull(source.getProperty("myBoolean").getAccessor());
       Assert.assertNotNull(source.getProperty("myBoolean").getMutator());
    }
+
+   @Test
+   public void testChangePropertyType()
+   {
+      JavaClassSource source = Roaster.create(JavaClassSource.class).setName("MyClass");
+      PropertySource<JavaClassSource> property = source.addProperty(Date.class, "myDate");
+      property.setType(Timestamp.class);
+      Assert.assertEquals("Timestamp", source.getField("myDate").getType().getName());
+      Assert.assertEquals("Timestamp", source.getMethod("getMyDate").getReturnType().getName());
+      Assert.assertNotNull(source.getMethod("setMyDate", Timestamp.class));
+
+   }
 }