Json writer and reader for enum with generic value (#221)

Tinkoff · Jul 26, 2023 · d732193 · d732193
1 parent f11639e
commit d732193
Show file tree

Hide file tree

Showing 22 changed files with 409 additions and 141 deletions.
diff --git a/...r/src/main/java/ru/tinkoff/kora/json/annotation/processor/reader/EnumReaderGenerator.java b/...r/src/main/java/ru/tinkoff/kora/json/annotation/processor/reader/EnumReaderGenerator.java
@@ -1,11 +1,14 @@
 package ru.tinkoff.kora.json.annotation.processor.reader;
 
 import com.squareup.javapoet.*;
+import ru.tinkoff.kora.annotation.processor.common.AnnotationUtils;
 import ru.tinkoff.kora.annotation.processor.common.CommonClassNames;
 import ru.tinkoff.kora.json.annotation.processor.JsonTypes;
 import ru.tinkoff.kora.json.annotation.processor.JsonUtils;
 
 import javax.annotation.Nullable;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import java.io.IOException;
@@ -14,6 +17,7 @@ public class EnumReaderGenerator {
 
     public TypeSpec generateForEnum(TypeElement typeElement) {
         var typeName = ClassName.get(typeElement);
+        var enumValue = this.detectValueType(typeElement);
 
         var typeBuilder = TypeSpec.classBuilder(JsonUtils.jsonReaderName(typeElement))
             .addAnnotation(AnnotationSpec.builder(CommonClassNames.koraGenerated)
@@ -22,13 +26,13 @@ public TypeSpec generateForEnum(TypeElement typeElement) {
             .addSuperinterface(ParameterizedTypeName.get(JsonTypes.jsonReader, typeName))
             .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
             .addOriginatingElement(typeElement);
-        var delegateType = ParameterizedTypeName.get(JsonTypes.enumJsonReader, typeName);
+        var delegateType = ParameterizedTypeName.get(JsonTypes.enumJsonReader, typeName, enumValue.type.box());
 
         typeBuilder.addField(delegateType, "delegate", Modifier.PRIVATE, Modifier.FINAL);
         typeBuilder.addMethod(MethodSpec.constructorBuilder()
             .addModifiers(Modifier.PUBLIC)
-            // todo detect string representation method for enum
-            .addCode("this.delegate = new $T<>($T.values(), v -> v.toString());\n", JsonTypes.enumJsonReader, typeName)
+            .addParameter(ParameterizedTypeName.get(JsonTypes.jsonReader, enumValue.type.box()), "valueReader")
+            .addCode("this.delegate = new $T<>($T.values(), $T::$N, valueReader);\n", JsonTypes.enumJsonReader, typeName, typeName, enumValue.accessor)
             .build());
         typeBuilder.addMethod(MethodSpec.methodBuilder("read")
             .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
@@ -42,4 +46,22 @@ public TypeSpec generateForEnum(TypeElement typeElement) {
         );
         return typeBuilder.build();
     }
+
+    record EnumValue(TypeName type, String accessor) {}
+
+    private EnumValue detectValueType(TypeElement typeElement) {
+        for (var enclosedElement : typeElement.getEnclosedElements()) {
+            if (!enclosedElement.getModifiers().contains(Modifier.PUBLIC)) continue;
+            if (enclosedElement.getModifiers().contains(Modifier.STATIC)) continue;
+            if (enclosedElement.getKind() != ElementKind.METHOD) continue;
+            if (enclosedElement instanceof ExecutableElement executableElement && executableElement.getParameters().isEmpty()) {
+                if (AnnotationUtils.isAnnotationPresent(executableElement, JsonTypes.json)) {
+                    var typeName = TypeName.get(executableElement.getReturnType());
+                    return new EnumValue(typeName, executableElement.getSimpleName().toString());
+                }
+            }
+        }
+        var typeName = ClassName.get(String.class);
+        return new EnumValue(typeName, "toString");
+    }
 }
diff --git a/...r/src/main/java/ru/tinkoff/kora/json/annotation/processor/writer/EnumWriterGenerator.java b/...r/src/main/java/ru/tinkoff/kora/json/annotation/processor/writer/EnumWriterGenerator.java
@@ -1,11 +1,14 @@
 package ru.tinkoff.kora.json.annotation.processor.writer;
 
 import com.squareup.javapoet.*;
+import ru.tinkoff.kora.annotation.processor.common.AnnotationUtils;
 import ru.tinkoff.kora.annotation.processor.common.CommonClassNames;
 import ru.tinkoff.kora.json.annotation.processor.JsonTypes;
 import ru.tinkoff.kora.json.annotation.processor.JsonUtils;
 
 import javax.annotation.Nullable;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import java.io.IOException;
@@ -21,12 +24,14 @@ public TypeSpec generateEnumWriter(TypeElement typeElement) {
             .addSuperinterface(ParameterizedTypeName.get(JsonTypes.jsonWriter, typeName))
             .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
             .addOriginatingElement(typeElement);
-        var delegateType = ParameterizedTypeName.get(JsonTypes.enumJsonWriter, typeName);
+        var enumValue = this.detectValueType(typeElement);
+        var delegateType = ParameterizedTypeName.get(JsonTypes.enumJsonWriter, typeName, enumValue.type.box());
 
         typeBuilder.addField(delegateType, "delegate", Modifier.PRIVATE, Modifier.FINAL);
         typeBuilder.addMethod(MethodSpec.constructorBuilder()
             .addModifiers(Modifier.PUBLIC)
-            .addCode("this.delegate = new $T<>($T.values(), v -> v.toString());\n", JsonTypes.enumJsonWriter, typeName)
+            .addParameter(ParameterizedTypeName.get(JsonTypes.jsonWriter, enumValue.type.box()), "valueWriter")
+            .addCode("this.delegate = new $T<>($T.values(), $T::$N, valueWriter);\n", JsonTypes.enumJsonWriter, typeName, typeName, enumValue.accessor)
             .build());
         typeBuilder.addMethod(MethodSpec.methodBuilder("write")
             .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
@@ -39,4 +44,22 @@ public TypeSpec generateEnumWriter(TypeElement typeElement) {
         );
         return typeBuilder.build();
     }
+
+    record EnumValue(TypeName type, String accessor) {}
+
+    private EnumValue detectValueType(TypeElement typeElement) {
+        for (var enclosedElement : typeElement.getEnclosedElements()) {
+            if (!enclosedElement.getModifiers().contains(Modifier.PUBLIC)) continue;
+            if (enclosedElement.getModifiers().contains(Modifier.STATIC)) continue;
+            if (enclosedElement.getKind() != ElementKind.METHOD) continue;
+            if (enclosedElement instanceof ExecutableElement executableElement && executableElement.getParameters().isEmpty()) {
+                if (AnnotationUtils.isAnnotationPresent(executableElement, JsonTypes.json)) {
+                    var typeName = TypeName.get(executableElement.getReturnType());
+                    return new EnumValue(typeName, executableElement.getSimpleName().toString());
+                }
+            }
+        }
+        var typeName = ClassName.get(String.class);
+        return new EnumValue(typeName, "toString");
+    }
 }
diff --git a/...nnotation-processor/src/test/java/ru/tinkoff/kora/json/annotation/processor/EnumTest.java b/...nnotation-processor/src/test/java/ru/tinkoff/kora/json/annotation/processor/EnumTest.java
@@ -1,13 +1,20 @@
 package ru.tinkoff.kora.json.annotation.processor;
 
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.JsonParser;
 import org.junit.jupiter.api.Test;
+import ru.tinkoff.kora.json.common.JsonReader;
+import ru.tinkoff.kora.json.common.JsonWriter;
 import ru.tinkoff.kora.kora.app.annotation.processor.KoraAppProcessor;
 
 import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class EnumTest extends AbstractJsonAnnotationProcessorTest {
+    JsonReader<String> stringReader = JsonParser::getValueAsString;
+    JsonWriter<String> stringWriter = JsonGenerator::writeString;
+
     @Test
     public void testEnum() {
         compile("""
@@ -19,11 +26,34 @@ public enum TestEnum {
 
         compileResult.assertSuccess();
 
-        var mapper = mapper("TestEnum");
+        var mapper = mapper("TestEnum", List.of(stringReader), List.of(stringWriter));
         mapper.verify(enumConstant("TestEnum", "VALUE1"), "\"VALUE1\"");
         mapper.verify(enumConstant("TestEnum", "VALUE2"), "\"VALUE2\"");
     }
 
+    @Test
+    public void testEnumWithCustomJsonValue() {
+        compile("""
+            @Json
+            public enum TestEnum {
+              VALUE1, VALUE2;
+              
+              @Json
+              public int intValue() {
+                return ordinal();
+              }
+            }
+            """);
+
+        compileResult.assertSuccess();
+        JsonReader<Integer> intReader = JsonParser::getIntValue;
+        JsonWriter<Integer> intWriter = JsonGenerator::writeNumber;
+
+        var mapper = mapper("TestEnum", List.of(intReader), List.of(intWriter));
+        mapper.verify(enumConstant("TestEnum", "VALUE1"), "0");
+        mapper.verify(enumConstant("TestEnum", "VALUE2"), "1");
+    }
+
 
     @Test
     public void testReaderFromExtension() {
@@ -34,13 +64,16 @@ enum TestEnum {
                 VALUE1, VALUE2
               }
               
+              default ru.tinkoff.kora.json.common.JsonReader<String> stringReader() { return com.fasterxml.jackson.core.JsonParser::getValueAsString; }
+              default ru.tinkoff.kora.json.common.JsonWriter<String> stringWriter() { return com.fasterxml.jackson.core.JsonGenerator::writeString; }
+
               @Root
               default String root(ru.tinkoff.kora.json.common.JsonReader<TestEnum> r) {return "";}
             }
             """);
 
         compileResult.assertSuccess();
-        assertThat(reader("TestApp_TestEnum")).isNotNull();
+        assertThat(reader("TestApp_TestEnum", stringReader)).isNotNull();
     }
 
     @Test
@@ -51,14 +84,17 @@ public interface TestApp {
               enum TestEnum {
                 VALUE1, VALUE2
               }
-              
+
+              default ru.tinkoff.kora.json.common.JsonReader<String> stringReader() { return com.fasterxml.jackson.core.JsonParser::getValueAsString; }
+              default ru.tinkoff.kora.json.common.JsonWriter<String> stringWriter() { return com.fasterxml.jackson.core.JsonGenerator::writeString; }
+
               @Root
               default String root(ru.tinkoff.kora.json.common.JsonWriter<TestEnum> r) {return "";}
             }
             """);
 
         compileResult.assertSuccess();
-        assertThat(writer("TestApp_TestEnum")).isNotNull();
+        assertThat(writer("TestApp_TestEnum", stringWriter)).isNotNull();
     }
 
     @Test
@@ -70,14 +106,17 @@ public interface TestApp {
               enum TestEnum {
                 VALUE1, VALUE2
               }
-              
+
+              default ru.tinkoff.kora.json.common.JsonReader<String> stringReader() { return com.fasterxml.jackson.core.JsonParser::getValueAsString; }
+              default ru.tinkoff.kora.json.common.JsonWriter<String> stringWriter() { return com.fasterxml.jackson.core.JsonGenerator::writeString; }
+
               @Root
               default String root(ru.tinkoff.kora.json.common.JsonReader<TestEnum> r) {return "";}
             }
             """);
 
         compileResult.assertSuccess();
-        assertThat(reader("TestApp_TestEnum")).isNotNull();
+        assertThat(reader("TestApp_TestEnum", stringReader)).isNotNull();
     }
 
     @Test
@@ -89,13 +128,16 @@ public interface TestApp {
               enum TestEnum {
                 VALUE1, VALUE2
               }
-              
+
+              default ru.tinkoff.kora.json.common.JsonReader<String> stringReader() { return com.fasterxml.jackson.core.JsonParser::getValueAsString; }
+              default ru.tinkoff.kora.json.common.JsonWriter<String> stringWriter() { return com.fasterxml.jackson.core.JsonGenerator::writeString; }
+
               @Root
               default String root(ru.tinkoff.kora.json.common.JsonWriter<TestEnum> r) {return "";}
             }
             """);
 
         compileResult.assertSuccess();
-        assertThat(writer("TestApp_TestEnum")).isNotNull();
+        assertThat(writer("TestApp_TestEnum", stringWriter)).isNotNull();
     }
 }
diff --git a/.../src/test/java/ru/tinkoff/kora/json/annotation/processor/JsonAnnotationProcessorTest.java b/.../src/test/java/ru/tinkoff/kora/json/annotation/processor/JsonAnnotationProcessorTest.java
@@ -321,22 +321,6 @@ void testNullableBeans() throws Exception {
             .hasMessageStartingWith("Expecting [VALUE_NUMBER_INT] token for field 'field4', got VALUE_NULL");
     }
 
-    @Test
-    void testEnum() throws Exception {
-        var cl = processClass0(DtoWithEnum.class);
-        var reader = cl.reader(DtoWithEnum.class, new EnumJsonReader<>(DtoWithEnum.TestEnum.values(), Enum::name));
-        var writer = cl.writer(DtoWithEnum.class, cl.writer(DtoWithEnum.TestEnum.class));
-
-        var expected = new DtoWithEnum(DtoWithEnum.TestEnum.VAL1);
-        var json = """
-            {
-              "testEnum" : "VAL1"
-            }""";
-
-        assertThat(fromJson(reader, json)).isEqualTo(expected);
-        assertThat(toJson(writer, expected)).isEqualTo(json);
-    }
-
     @Test
     void testObject() throws Exception {
         var cl = processClass0(DtoWithObject.class);

diff --git a/...on-processor/src/test/java/ru/tinkoff/kora/json/annotation/processor/dto/DtoWithEnum.java b/...on-processor/src/test/java/ru/tinkoff/kora/json/annotation/processor/dto/DtoWithEnum.java
diff --git a/json/json-common/src/main/java/ru/tinkoff/kora/json/common/EnumJsonReader.java b/json/json-common/src/main/java/ru/tinkoff/kora/json/common/EnumJsonReader.java
@@ -2,38 +2,35 @@
 
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.JsonParser;
-import com.fasterxml.jackson.core.JsonToken;
 
 import javax.annotation.Nullable;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
 
-public final class EnumJsonReader<T extends Enum<T>> implements JsonReader<T> {
-    private final Map<String, T> values;
+public final class EnumJsonReader<T extends Enum<T>, V> implements JsonReader<T> {
+    private final Map<V, T> values;
+    private final JsonReader<V> valueReader;
 
-    public EnumJsonReader(T[] values, Function<T, String> mapper) {
+    public EnumJsonReader(T[] values, Function<T, V> mapper, JsonReader<V> valueReader) {
         this.values = new HashMap<>();
         for (var value : values) {
             this.values.put(mapper.apply(value), value);
         }
+        this.valueReader = valueReader;
     }
 
     @Nullable
     @Override
     public T read(JsonParser parser) throws IOException {
-        var token = parser.currentToken();
-        if (token == JsonToken.VALUE_NULL) {
+        var jsonValue = this.valueReader.read(parser);
+        if (jsonValue == null) {
             return null;
         }
-        if (token != JsonToken.VALUE_STRING) {
-            throw new JsonParseException(parser, "Expecting VALUE_STRING token, got " + token);
-        }
-        var stringValue = parser.getText();
-        var value = this.values.get(stringValue);
+        var value = this.values.get(jsonValue);
         if (value == null) {
-            throw new JsonParseException(parser, "Expecting one of " + this.values.keySet() + ", got " + stringValue);
+            throw new JsonParseException(parser, "Expecting one of " + this.values.keySet() + ", got " + jsonValue);
         }
         return value;
     }

diff --git a/json/json-common/src/main/java/ru/tinkoff/kora/json/common/EnumJsonWriter.java b/json/json-common/src/main/java/ru/tinkoff/kora/json/common/EnumJsonWriter.java
@@ -1,28 +1,35 @@
 package ru.tinkoff.kora.json.common;
 
 import com.fasterxml.jackson.core.JsonGenerator;
-import com.fasterxml.jackson.core.io.SerializedString;
 
 import javax.annotation.Nullable;
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.util.function.Function;
 
-public final class EnumJsonWriter<T extends Enum<T>> implements JsonWriter<T> {
-    private final SerializedString[] values;
+public final class EnumJsonWriter<T extends Enum<T>, V> implements JsonWriter<T> {
+    private final RawJson[] values;
 
-    public EnumJsonWriter(T[] values, Function<T, String> mapper) {
-        this.values = new SerializedString[values.length];
+    public EnumJsonWriter(T[] values, Function<T, V> valueExtractor, JsonWriter<V> valueWriter) {
+        this.values = new RawJson[values.length];
         for (int i = 0; i < values.length; i++) {
-            this.values[i] = new SerializedString(mapper.apply(values[i]));
+            var enumValue = values[i];
+            var value = valueExtractor.apply(enumValue);
+            try {
+                var bytes = valueWriter.toByteArray(value);
+                this.values[i] = new RawJson(bytes);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
         }
     }
 
     @Override
     public void write(JsonGenerator gen, @Nullable T object) throws IOException {
         if (object == null) {
             gen.writeNull();
-        } else {
-            gen.writeString(this.values[object.ordinal()]);
+            return;
         }
+        gen.writeRawValue(this.values[object.ordinal()]);
     }
 }