Merge pull request #458 from blizznets/master

#455 Override builder methods from parent class
joelittlejohn · Dec 3, 2015 · 1338fee · 1338fee
2 parents 4964303 + a54409e
commit 1338fee
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 6 deletions.
diff --git a/jsonschema2pojo-core/src/main/java/org/jsonschema2pojo/rules/PropertiesRule.java b/jsonschema2pojo-core/src/main/java/org/jsonschema2pojo/rules/PropertiesRule.java
@@ -16,15 +16,15 @@
 
 package org.jsonschema2pojo.rules;
 
-import java.util.Iterator;
-
 import com.fasterxml.jackson.databind.JsonNode;
+import com.sun.codemodel.*;
 import org.jsonschema2pojo.Schema;
-import com.sun.codemodel.JDefinedClass;
+
+import java.util.Iterator;
 
 /**
  * Applies the "properties" schema rule.
- * 
+ *
  * @see <a
  *      href="http://tools.ietf.org/html/draft-zyp-json-schema-03#section-5.2">http://tools.ietf.org/html/draft-zyp-json-schema-03#section-5.2</a>
  */
@@ -41,7 +41,7 @@ protected PropertiesRule(RuleFactory ruleFactory) {
      * <p>
      * For each property present within the properties node, this rule will
      * invoke the 'property' rule provided by the given schema mapper.
-     * 
+     *
      * @param nodeName
      *            the name of the node for which properties are being added
      * @param node
@@ -54,14 +54,44 @@ protected PropertiesRule(RuleFactory ruleFactory) {
     @Override
     public JDefinedClass apply(String nodeName, JsonNode node, JDefinedClass jclass, Schema schema) {
 
-        for (Iterator<String> properties = node.fieldNames(); properties.hasNext();) {
+        for (Iterator<String> properties = node.fieldNames(); properties.hasNext(); ) {
             String property = properties.next();
 
             ruleFactory.getPropertyRule().apply(property, node.get(property), jclass, schema);
         }
 
+        if (ruleFactory.getGenerationConfig().isGenerateBuilders()) {
+            if (!jclass._extends().name().equals("Object")) {
+                addOverrideBuilders(jclass, jclass.owner()._getClass(jclass._extends().fullName()));
+            }
+        }
+
         ruleFactory.getAnnotator().propertyOrder(jclass, node);
 
         return jclass;
     }
+
+    private void addOverrideBuilders(JDefinedClass jclass, JDefinedClass parentJclass) {
+        if (parentJclass == null) {
+            return;
+        }
+
+        for (JMethod parentJMethod : parentJclass.methods()) {
+            if (parentJMethod.name().startsWith("with") && parentJMethod.params().size() == 1) {
+                addOverrideBuilder(jclass, parentJMethod, parentJMethod.params().get(0));
+            }
+        }
+    }
+
+    private JMethod addOverrideBuilder(JDefinedClass thisJDefinedClass, JMethod parentBuilder, JVar parentParam) {
+        JMethod builder = thisJDefinedClass.method(parentBuilder.mods().getValue(), thisJDefinedClass, parentBuilder.name());
+        builder.annotate(Override.class);
+
+        JVar param = builder.param(parentParam.type(), parentParam.name());
+        JBlock body = builder.body();
+        body.invoke(JExpr._super(), parentBuilder).arg(param);
+        body._return(JExpr._this());
+
+        return builder;
+    }
 }
diff --git a/...chema2pojo-integration-tests/src/test/java/org/jsonschema2pojo/integration/ExtendsIT.java b/...chema2pojo-integration-tests/src/test/java/org/jsonschema2pojo/integration/ExtendsIT.java
@@ -13,6 +13,8 @@
 package org.jsonschema2pojo.integration;
 
 import java.beans.PropertyDescriptor;
+import java.lang.reflect.Method;
+
 import static org.jsonschema2pojo.integration.util.CodeGenerationHelper.*;
 import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
@@ -100,4 +102,21 @@ public void extendsSchemaWithinDefinitions() throws Exception {
 
         assertThat(subtype.getSuperclass(), is(equalTo(supertype)));
     }
+
+    @Test
+    @SuppressWarnings("rawtypes")
+    public void extendsBuilderMethods() throws Exception {
+        ClassLoader resultsClassLoader = generateAndCompile("/schema/extends/subtypeOfSubtypeOfA.json", "com.example", config("generateBuilders", true));
+
+        Class subtype = resultsClassLoader.loadClass("com.example.SubtypeOfSubtypeOfA");
+        Class supertype = resultsClassLoader.loadClass("com.example.SubtypeOfSubtypeOfAParent");
+
+        Method builderMethod = supertype.getDeclaredMethod("withParent", String.class);
+        assertNotNull("no withParent method", builderMethod);
+        assertThat(builderMethod.getReturnType(), is(equalTo(supertype)));
+
+        Method builderMethodOverride = subtype.getDeclaredMethod("withParent", String.class);
+        assertNotNull("no withParent method", builderMethodOverride);
+        assertThat(builderMethodOverride.getReturnType(), is(equalTo(subtype)));
+    }
 }