Skip to content

Commit 4e589d7

Browse files
erdichenErdi Chen
andauthored
Forward Thrift annotations to TypeProcessor as tags (#26)
Add Thrift annotations as a tag in the type specs passed to the `TypeProcessor` post processor. Supports the generated data classes, enum classes, and sealed classes. Access the annotations in the post processor: 1. `propertySpec.tag<ThriftAnnotations>()?.annotations` 2. `typeSpec.tag<ThriftAnnotations>()?.annotations` Co-authored-by: Erdi Chen <erdic@x.com>
1 parent d2d89aa commit 4e589d7

File tree

3 files changed

+269
-0
lines changed

3 files changed

+269
-0
lines changed
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Thrifty
3+
*
4+
* Copyright (c) Benjamin Bader
5+
* Copyright (c) Microsoft Corporation
6+
*
7+
* All rights reserved.
8+
*
9+
* Licensed under the Apache License, Version 2.0 (the License);
10+
* you may not use this file except in compliance with the License.
11+
* You may obtain a copy of the License at
12+
*
13+
* http://www.apache.org/licenses/LICENSE-2.0
14+
*
15+
* THIS CODE IS PROVIDED ON AN *AS IS* BASIS, WITHOUT WARRANTIES OR
16+
* CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
17+
* WITHOUT LIMITATION ANY IMPLIED WARRANTIES OR CONDITIONS OF TITLE,
18+
* FITNESS FOR A PARTICULAR PURPOSE, MERCHANTABLITY OR NON-INFRINGEMENT.
19+
*
20+
* See the Apache Version 2.0 License for specific language governing permissions and limitations under the License.
21+
*/
22+
package com.bendb.thrifty.compiler.spi;
23+
24+
import java.util.Map;
25+
import java.util.Objects;
26+
27+
public class ThriftAnnotations {
28+
private final Map<String, String> annotations;
29+
30+
public ThriftAnnotations(Map<String, String> annotations) {
31+
this.annotations = annotations;
32+
}
33+
34+
public Map<String, String> getAnnotations() {
35+
return annotations;
36+
}
37+
38+
@Override
39+
public boolean equals(Object o) {
40+
if (this == o) return true;
41+
if (o == null || getClass() != o.getClass()) return false;
42+
ThriftAnnotations that = (ThriftAnnotations) o;
43+
return Objects.equals(annotations, that.annotations);
44+
}
45+
46+
@Override
47+
public int hashCode() {
48+
return Objects.hash(annotations);
49+
}
50+
}

thrifty-kotlin-codegen/src/main/kotlin/com/bendb/thrifty/kgen/KotlinCodeGenerator.kt

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import com.bendb.thrifty.Struct
2828
import com.bendb.thrifty.TType
2929
import com.bendb.thrifty.ThriftException
3030
import com.bendb.thrifty.ThriftField
31+
import com.bendb.thrifty.compiler.spi.ThriftAnnotations
3132
import com.bendb.thrifty.compiler.spi.TypeProcessor
3233
import com.bendb.thrifty.protocol.MessageMetadata
3334
import com.bendb.thrifty.protocol.Protocol
@@ -80,6 +81,7 @@ import com.squareup.kotlinpoet.NameAllocator
8081
import com.squareup.kotlinpoet.ParameterSpec
8182
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
8283
import com.squareup.kotlinpoet.PropertySpec
84+
import com.squareup.kotlinpoet.Taggable
8385
import com.squareup.kotlinpoet.TypeAliasSpec
8486
import com.squareup.kotlinpoet.TypeName
8587
import com.squareup.kotlinpoet.TypeSpec
@@ -124,6 +126,14 @@ private object ClassNames {
124126
val IO_EXCEPTION = ClassName("okio", "IOException")
125127
}
126128

129+
private fun <T : Taggable.Builder<T>> Taggable.Builder<T>.tagWithAnnotations(
130+
annotations: Map<String, String>?
131+
) {
132+
if (annotations?.isNotEmpty() == true) {
133+
tag(ThriftAnnotations::class, ThriftAnnotations(annotations))
134+
}
135+
}
136+
127137
/**
128138
* Generates Kotlin code from a [Schema].
129139
*
@@ -424,11 +434,15 @@ class KotlinCodeGenerator(fieldNamingPolicy: FieldNamingPolicy = FieldNamingPoli
424434
if (member.isDeprecated) enumMemberSpec.addAnnotation(makeDeprecated())
425435
if (member.hasJavadoc) enumMemberSpec.addKdoc("%L", member.documentation)
426436

437+
enumMemberSpec.tagWithAnnotations(member.annotations)
438+
427439
val name = nameAllocator.get(member)
428440
typeBuilder.addEnumConstant(name, enumMemberSpec.build())
429441
findByValue.addStatement("%L -> %L", member.value, name)
430442
}
431443

444+
typeBuilder.tagWithAnnotations(enumType.annotations)
445+
432446
val companion =
433447
TypeSpec.companionObjectBuilder()
434448
.addFunction(findByValue.addStatement("else -> null").endControlFlow().build())
@@ -517,6 +531,8 @@ class KotlinCodeGenerator(fieldNamingPolicy: FieldNamingPolicy = FieldNamingPoli
517531
.jvmField()
518532
.addAnnotation(thriftField)
519533

534+
prop.tagWithAnnotations(field.annotations)
535+
520536
if (field.hasJavadoc) prop.addKdoc("%L", field.documentation)
521537
if (field.isObfuscated) prop.addAnnotation(Obfuscated::class)
522538
if (field.isRedacted) prop.addAnnotation(Redacted::class)
@@ -525,6 +541,8 @@ class KotlinCodeGenerator(fieldNamingPolicy: FieldNamingPolicy = FieldNamingPoli
525541
typeBuilder.addProperty(prop.build())
526542
}
527543

544+
typeBuilder.tagWithAnnotations(struct.annotations)
545+
528546
val adapterTypeName = ClassName(struct.kotlinNamespace, struct.name, "${struct.name}Adapter")
529547
val adapterInterfaceTypeName = Adapter::class.asTypeName().parameterizedBy(struct.typeName)
530548

@@ -588,6 +606,8 @@ class KotlinCodeGenerator(fieldNamingPolicy: FieldNamingPolicy = FieldNamingPoli
588606
if (struct.hasJavadoc) addKdoc("%L", struct.documentation)
589607
}
590608

609+
typeBuilder.tagWithAnnotations(struct.annotations)
610+
591611
var defaultValueTypeName: ClassName? = null
592612
val nameAllocator = nameAllocators[struct]
593613
for (field in struct.fields) {
@@ -608,6 +628,8 @@ class KotlinCodeGenerator(fieldNamingPolicy: FieldNamingPolicy = FieldNamingPoli
608628
.primaryConstructor(propConstructor.build())
609629
.tag(FieldIdMarker(field.id))
610630

631+
dataPropBuilder.tagWithAnnotations(field.annotations)
632+
611633
val toStringBody = buildCodeBlock {
612634
add("return \"${struct.name}($name=")
613635
when {
Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
/*
2+
* Thrifty
3+
*
4+
* Copyright (c) Benjamin Bader
5+
* Copyright (c) Microsoft Corporation
6+
*
7+
* All rights reserved.
8+
*
9+
* Licensed under the Apache License, Version 2.0 (the License);
10+
* you may not use this file except in compliance with the License.
11+
* You may obtain a copy of the License at
12+
*
13+
* http://www.apache.org/licenses/LICENSE-2.0
14+
*
15+
* THIS CODE IS PROVIDED ON AN *AS IS* BASIS, WITHOUT WARRANTIES OR
16+
* CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
17+
* WITHOUT LIMITATION ANY IMPLIED WARRANTIES OR CONDITIONS OF TITLE,
18+
* FITNESS FOR A PARTICULAR PURPOSE, MERCHANTABLITY OR NON-INFRINGEMENT.
19+
*
20+
* See the Apache Version 2.0 License for specific language governing permissions and limitations under the License.
21+
*/
22+
package com.bendb.thrifty.kgen
23+
24+
import com.bendb.thrifty.compiler.spi.ThriftAnnotations
25+
import com.bendb.thrifty.compiler.spi.TypeProcessor
26+
import com.bendb.thrifty.schema.FieldNamingPolicy
27+
import com.bendb.thrifty.schema.Loader
28+
import com.bendb.thrifty.schema.Schema
29+
import com.squareup.kotlinpoet.AnnotationSpec
30+
import com.squareup.kotlinpoet.DelicateKotlinPoetApi
31+
import com.squareup.kotlinpoet.FileSpec
32+
import com.squareup.kotlinpoet.TypeSpec
33+
import com.squareup.kotlinpoet.tag
34+
import io.kotest.matchers.string.shouldContain
35+
import java.io.File
36+
import org.junit.jupiter.api.Test
37+
import org.junit.jupiter.api.io.TempDir
38+
import org.junit.jupiter.api.parallel.Execution
39+
import org.junit.jupiter.api.parallel.ExecutionMode
40+
41+
@Execution(ExecutionMode.CONCURRENT)
42+
class ThriftAnnotationTests {
43+
@TempDir lateinit var tempDir: File
44+
45+
annotation class TestAnnotation(val name: String, val value: String)
46+
47+
class ThriftAnnotationKotlinProcessor : TypeProcessor {
48+
@OptIn(DelicateKotlinPoetApi::class)
49+
override fun process(type: TypeSpec): TypeSpec {
50+
val newEnumConstants =
51+
type.enumConstants.entries.map { (name, value) -> name to process(value) }
52+
val newPropertySpecs =
53+
type.propertySpecs.map { prop ->
54+
val annotations = prop.tag<ThriftAnnotations>()?.annotations ?: return@map prop
55+
val newAnnotations =
56+
annotations.map { (name, value) ->
57+
AnnotationSpec.builder(TestAnnotation::class.java)
58+
.addMember("%L = \"%L\"", name, value)
59+
.build()
60+
}
61+
return@map prop.toBuilder().addAnnotations(newAnnotations).build()
62+
}
63+
val newTypeSpecs = type.typeSpecs.map { spec -> process(spec) }
64+
return type.toBuilder().run {
65+
type.tag<ThriftAnnotations>()?.annotations?.also { annotations ->
66+
val newAnnotations =
67+
annotations.map { (name, value) ->
68+
AnnotationSpec.builder(TestAnnotation::class.java)
69+
.addMember("%L = \"%L\"", name, value)
70+
.build()
71+
}
72+
addAnnotations(newAnnotations)
73+
}
74+
enumConstants.clear()
75+
enumConstants.putAll(newEnumConstants)
76+
propertySpecs.clear()
77+
propertySpecs.addAll(newPropertySpecs)
78+
typeSpecs.clear()
79+
typeSpecs.addAll(newTypeSpecs)
80+
build()
81+
}
82+
}
83+
}
84+
85+
@Test
86+
fun `struct to data class should generate annotation from custom Thrift annotation`() {
87+
val thrift =
88+
"""
89+
namespace kt com.test
90+
91+
struct Test {
92+
1: optional string jsonName (json.name = "json_name");
93+
2: optional string kotlin_name (go.name = "GoName", kotlin.name = "kotlinName");
94+
} (kotlin.name = "KotlinTest")
95+
"""
96+
.trimIndent()
97+
98+
val file = generate(thrift).single()
99+
100+
file.toString().also {
101+
it shouldContain
102+
"""
103+
@ThriftAnnotationTests.TestAnnotation(kotlin.name = "KotlinTest")
104+
"""
105+
.trimIndent()
106+
it shouldContain
107+
"""
108+
@ThriftAnnotationTests.TestAnnotation(json.name = "json_name")
109+
"""
110+
.trimIndent()
111+
it shouldContain
112+
"""
113+
@ThriftAnnotationTests.TestAnnotation(go.name = "GoName")
114+
"""
115+
.trimIndent()
116+
it shouldContain
117+
"""
118+
@ThriftAnnotationTests.TestAnnotation(kotlin.name = "kotlinName")
119+
"""
120+
.trimIndent()
121+
}
122+
}
123+
124+
@Test
125+
fun `union generate sealed class should generate annotation from custom Thrift annotation`() {
126+
val thrift =
127+
"""
128+
namespace kt com.test
129+
130+
union Union {
131+
1: i32 Foo;
132+
2: i64 Bar;
133+
3: string Baz;
134+
4: i32 NotFoo (java.name = "NotFoo");
135+
} (kotlin.name = "UnionClass")
136+
"""
137+
.trimMargin()
138+
139+
val file = generate(thrift)
140+
141+
file.single().toString().also {
142+
it shouldContain
143+
"""
144+
@ThriftAnnotationTests.TestAnnotation(kotlin.name = "UnionClass")
145+
"""
146+
.trimIndent()
147+
it shouldContain
148+
"""
149+
@ThriftAnnotationTests.TestAnnotation(java.name = "NotFoo")
150+
"""
151+
.trimIndent()
152+
}
153+
}
154+
155+
@Test
156+
fun `enum generation should generate annotation from custom Thrift annotation`() {
157+
val thrift =
158+
"""
159+
namespace kt com.test
160+
161+
enum Foo {
162+
FIRST_VALUE = 0 (go.name = "FirstValue"),
163+
SECOND_VALUE = 1,
164+
THIRD_VALUE = 2
165+
} (kotlin.name = "FooEnum")
166+
"""
167+
.trimIndent()
168+
169+
val file = generate(thrift)
170+
171+
file.single().toString().also {
172+
it shouldContain
173+
"""
174+
@ThriftAnnotationTests.TestAnnotation(go.name = "FirstValue")
175+
"""
176+
.trimIndent()
177+
it shouldContain
178+
"""
179+
@ThriftAnnotationTests.TestAnnotation(kotlin.name = "FooEnum")
180+
"""
181+
.trimIndent()
182+
}
183+
}
184+
185+
private fun generate(thrift: String): List<FileSpec> {
186+
187+
return KotlinCodeGenerator(FieldNamingPolicy.JAVA)
188+
.apply { processor = ThriftAnnotationKotlinProcessor() }
189+
.generate(load(thrift))
190+
}
191+
192+
private fun load(thrift: String): Schema {
193+
val file = File(tempDir, "test.thrift").also { it.writeText(thrift) }
194+
val loader = Loader().apply { addThriftFile(file.toPath()) }
195+
return loader.load()
196+
}
197+
}

0 commit comments

Comments
 (0)