feat: add support for JsonSubTypes (#1336)

packages/java/engine-core/pom.xml

@@ -73,6 +73,11 @@
       <artifactId>parser-jvm-plugin-nonnull</artifactId>
       <version>${project.version}</version>
     </dependency>
+    <dependency>
+      <groupId>dev.hilla</groupId>
+      <artifactId>parser-jvm-plugin-subtypes</artifactId>
+      <version>${project.version}</version>
+    </dependency>
     <dependency>
       <groupId>dev.hilla</groupId>
       <artifactId>parser-jvm-plugin-model</artifactId>

packages/java/engine-core/src/main/java/dev/hilla/engine/GeneratorConfiguration.java

@@ -142,7 +142,8 @@ static class PluginsProcessor extends ConfigList.Processor<Plugin> {
                 new Plugin("@hilla/generator-typescript-plugin-client"),
                 new Plugin("@hilla/generator-typescript-plugin-barrel"),
                 new Plugin("@hilla/generator-typescript-plugin-model"),
-                new Plugin("@hilla/generator-typescript-plugin-push"));
+                new Plugin("@hilla/generator-typescript-plugin-push"),
+                new Plugin("@hilla/generator-typescript-plugin-subtypes"));
 
         PluginsProcessor() {
             super(DEFAULTS);

packages/java/engine-core/src/main/java/dev/hilla/engine/GeneratorProcessor.java

@@ -25,6 +25,7 @@
 @NpmPackage(value = "@hilla/generator-typescript-plugin-barrel", version = "2.3.0-alpha5")
 @NpmPackage(value = "@hilla/generator-typescript-plugin-model", version = "2.3.0-alpha5")
 @NpmPackage(value = "@hilla/generator-typescript-plugin-push", version = "2.3.0-alpha5")
+@NpmPackage(value = "@hilla/generator-typescript-plugin-subtypes", version = "2.3.0-alpha5")
 public final class GeneratorProcessor {
     private static final Logger logger = LoggerFactory
             .getLogger(GeneratorProcessor.class);

packages/java/engine-core/src/main/java/dev/hilla/engine/ParserConfiguration.java

@@ -12,6 +12,7 @@
 import dev.hilla.parser.plugins.backbone.BackbonePlugin;
 import dev.hilla.parser.plugins.model.ModelPlugin;
 import dev.hilla.parser.plugins.nonnull.NonnullPlugin;
+import dev.hilla.parser.plugins.subtypes.SubTypesPlugin;
 import dev.hilla.parser.plugins.transfertypes.TransferTypesPlugin;
 import dev.hilla.parser.utils.ConfigList;
 
@@ -188,6 +189,7 @@ static class PluginsProcessor extends ConfigList.Processor<Plugin> {
                 new Plugin(BackbonePlugin.class.getName()),
                 new Plugin(TransferTypesPlugin.class.getName()),
                 new Plugin(NonnullPlugin.class.getName()),
+                new Plugin(SubTypesPlugin.class.getName()),
                 new Plugin(ModelPlugin.class.getName()));
 
         PluginsProcessor() {

packages/java/parser-jvm-plugin-backbone/src/main/java/dev/hilla/parser/plugins/backbone/EntityPlugin.java

@@ -67,8 +67,8 @@ public NodeDependencies scan(@Nonnull NodeDependencies nodeDependencies) {
                 Stream.of(EntityNode.of(ref.getClassInfo())));
     }
 
-    private void attachSchemaWithNameToOpenApi(Schema<?> schema, String name,
-            OpenAPI openApi) {
+    public static void attachSchemaWithNameToOpenApi(Schema<?> schema,
+            String name, OpenAPI openApi) {
         var components = openApi.getComponents();
 
         if (components == null) {

packages/java/parser-jvm-plugin-subtypes/pom.xml

@@ -0,0 +1,58 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xmlns="http://maven.apache.org/POM/4.0.0"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+
+  <parent>
+    <groupId>dev.hilla</groupId>
+    <artifactId>hilla-project</artifactId>
+    <version>2.3-SNAPSHOT</version>
+    <relativePath>../../../pom.xml</relativePath>
+  </parent>
+
+  <artifactId>parser-jvm-plugin-subtypes</artifactId>
+  <name>Hilla JVM Parser SubTypes Plugin</name>
+  <packaging>jar</packaging>
+
+  <properties>
+    <formatter.basedir>${project.parent.basedir}</formatter.basedir>
+  </properties>
+
+  <dependencies>
+    <dependency>
+      <groupId>dev.hilla</groupId>
+      <artifactId>parser-jvm-core</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>com.google.code.findbugs</groupId>
+      <artifactId>jsr305</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.github.classgraph</groupId>
+      <artifactId>classgraph</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.junit.jupiter</groupId>
+      <artifactId>junit-jupiter</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>dev.hilla</groupId>
+      <artifactId>parser-jvm-utils</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>dev.hilla</groupId>
+      <artifactId>parser-jvm-plugin-backbone</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>dev.hilla</groupId>
+      <artifactId>parser-jvm-test-utils</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+</project>

packages/java/parser-jvm-plugin-subtypes/src/main/java/dev/hilla/parser/plugins/subtypes/SubTypesPlugin.java

@@ -0,0 +1,155 @@
+package dev.hilla.parser.plugins.subtypes;
+
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import dev.hilla.parser.core.AbstractNode;
+import dev.hilla.parser.core.AbstractPlugin;
+import dev.hilla.parser.core.Node;
+import dev.hilla.parser.core.NodeDependencies;
+import dev.hilla.parser.core.NodePath;
+import dev.hilla.parser.core.Plugin;
+import dev.hilla.parser.core.PluginConfiguration;
+import dev.hilla.parser.models.ClassInfoModel;
+import dev.hilla.parser.models.ClassRefSignatureModel;
+import dev.hilla.parser.plugins.backbone.BackbonePlugin;
+import dev.hilla.parser.plugins.backbone.EntityPlugin;
+import dev.hilla.parser.plugins.backbone.nodes.EntityNode;
+import dev.hilla.parser.plugins.backbone.nodes.TypedNode;
+import io.swagger.v3.oas.models.OpenAPI;
+import io.swagger.v3.oas.models.media.ComposedSchema;
+import io.swagger.v3.oas.models.media.ObjectSchema;
+import io.swagger.v3.oas.models.media.Schema;
+import io.swagger.v3.oas.models.media.StringSchema;
+
+import javax.annotation.Nonnull;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+/**
+ * This plugin adds support for {@code @JsonTypeInfo} and {@code @JsonSubTypes}.
+ */
+public final class SubTypesPlugin extends AbstractPlugin<PluginConfiguration> {
+    @Override
+    public void enter(NodePath<?> nodePath) {
+    }
+
+    @Override
+    public void exit(NodePath<?> nodePath) {
+        // deal with the union nodes, which does not correspond to an existing
+        // class, but express the union of all the @JsonSubTypes
+        if (nodePath.getNode() instanceof UnionNode) {
+            var unionNode = (UnionNode) nodePath.getNode();
+            var cls = (Class<?>) unionNode.getSource().get();
+
+            // verify that the class has a @JsonTypeInfo annotation
+            // and then add all the @JsonSubTypes to the schema as a `oneOf`
+            if (cls.getAnnotationsByType(JsonTypeInfo.class).length > 0) {
+                var schema = (Schema<?>) unionNode.getTarget();
+                getJsonSubTypes(cls).map(JsonSubTypes.Type::value)
+                        .forEach(c -> {
+                            schema.addOneOfItem(new Schema<Object>() {
+                                {
+                                    set$ref("#/components/schemas/"
+                                            + c.getName());
+                                }
+                            });
+                        });
+            }
+
+            // attach the schema to the openapi
+            EntityPlugin.attachSchemaWithNameToOpenApi(unionNode.getTarget(),
+                    cls.getName() + "Union",
+                    (OpenAPI) nodePath.getParentPath().getNode().getTarget());
+        }
+
+        // entity nodes whose superclass has a @JsonSubTypes annotation must
+        // have a @type property whose value comes from the annotation
+        if (nodePath.getNode() instanceof EntityNode) {
+            var entityNode = (EntityNode) nodePath.getNode();
+            var cls = (Class<?>) entityNode.getSource().get();
+
+            Optional.ofNullable(cls.getSuperclass())
+                    .map(SubTypesPlugin::getJsonSubTypes).stream()
+                    .flatMap(Function.identity())
+                    .filter(t -> cls.equals(t.value())).findAny()
+                    .ifPresent(t -> {
+                        var schema = (ComposedSchema) entityNode.getTarget();
+                        schema.getAnyOf().stream()
+                                .filter(s -> s instanceof ObjectSchema)
+                                .map(ObjectSchema.class::cast)
+                                .forEach(s -> s.addProperty("@type",
+                                        new StringSchema() {
+                                            {
+                                                setType("string");
+                                                setExample(t.name());
+                                            }
+                                        }));
+                    });
+        }
+    }
+
+    @Override
+    public Collection<Class<? extends Plugin>> getRequiredPlugins() {
+        return List.of(BackbonePlugin.class);
+    }
+
+    @Nonnull
+    @Override
+    public NodeDependencies scan(@Nonnull NodeDependencies nodeDependencies) {
+        if (!(nodeDependencies.getNode() instanceof TypedNode)) {
+            return nodeDependencies;
+        }
+
+        var typedNode = (TypedNode) nodeDependencies.getNode();
+        if (!(typedNode.getType() instanceof ClassRefSignatureModel)) {
+            return nodeDependencies;
+        }
+
+        var ref = (ClassRefSignatureModel) typedNode.getType();
+        if (ref.isJDKClass() || ref.isDate() || ref.isIterable()) {
+            return nodeDependencies;
+        }
+
+        // all types mentioned in @JsonSubTypes must be parsed, even if they are
+        // not used directly
+        Class<?> refClass = (Class<?>) ref.getClassInfo().get();
+        var subTypes = getJsonSubTypes(refClass).map(JsonSubTypes.Type::value)
+                .map(ClassInfoModel::of).<Node<?, ?>> map(EntityNode::of);
+
+        // create a union node for classes annotated with @JsonTypeInfo
+        if (refClass.getAnnotationsByType(JsonTypeInfo.class).length > 0) {
+            var unionType = UnionNode.of(ref.getClassInfo());
+            subTypes = Stream.concat(Stream.of(unionType), subTypes);
+        }
+
+        return nodeDependencies.appendRelatedNodes(subTypes);
+    }
+
+    private static Stream<JsonSubTypes.Type> getJsonSubTypes(Class<?> cls) {
+        return Optional.of(cls)
+                .map(c -> c.getAnnotationsByType(JsonSubTypes.class))
+                .filter(a -> a.length > 0).map(a -> a[0])
+                .map(JsonSubTypes::value).stream().flatMap(Arrays::stream);
+    }
+
+    /**
+     * A node that represents the union of all the mentioned subclasses of a
+     * class annotated with {@code @JsonSubTypes}.
+     */
+    public static class UnionNode
+            extends AbstractNode<ClassInfoModel, Schema<?>> {
+        private UnionNode(@Nonnull ClassInfoModel source,
+                @Nonnull ObjectSchema target) {
+            super(source, target);
+        }
+
+        @Nonnull
+        static public UnionNode of(@Nonnull ClassInfoModel model) {
+            return new UnionNode(model, new ObjectSchema());
+        }
+    }
+}

packages/java/parser-jvm-plugin-subtypes/src/main/java/module-info.java

@@ -0,0 +1,12 @@
+module dev.hilla.parser.plugins.subtypes {
+    requires com.fasterxml.jackson.databind;
+    requires dev.hilla.parser.plugins.backbone;
+    requires dev.hilla.parser.utils;
+    requires jsr305;
+    requires jakarta.annotation;
+    requires io.swagger.v3.oas.models;
+    requires io.github.classgraph;
+    requires dev.hilla.parser.core;
+
+    exports dev.hilla.parser.plugins.subtypes;
+}

packages/java/parser-jvm-plugin-subtypes/src/test/java/dev/hilla/parser/plugins/subtypes/AddEvent.java

@@ -0,0 +1,6 @@
+package dev.hilla.parser.plugins.subtypes;
+
+public class AddEvent extends BaseEvent {
+
+    public String item;
+}

packages/java/parser-jvm-plugin-subtypes/src/test/java/dev/hilla/parser/plugins/subtypes/BaseEvent.java

@@ -0,0 +1,12 @@
+package dev.hilla.parser.plugins.subtypes;
+
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY)
+@JsonSubTypes({ @JsonSubTypes.Type(value = AddEvent.class, name = "add"),
+        @JsonSubTypes.Type(value = UpdateEvent.class, name = "update"),
+        @JsonSubTypes.Type(value = DeleteEvent.class, name = "delete") })
+public class BaseEvent {
+    public int id;
+}