Skip to content

Commit

Permalink
Add register service and generate description to file (#28)
Browse files Browse the repository at this point in the history
Co-authored-by: wz <[email protected]>
  • Loading branch information
GuoDuanLZ and GuoDuanLZ authored Jun 12, 2020
1 parent 2a4b99e commit cb9937b
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 47 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -29,3 +29,9 @@ fun String.toPlatformPath(): String {
}
}
}

fun String.replaceExtensionName(old: String, new: String): String {
val extension = ".$old"
if (!this.endsWith(extension)) return this
return this.substring(0, this.length - old.length - 1) + ".$new"
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import kotlin.reflect.full.companionObjectInstance

object ProtoTypes {
private val protoToClassMap: BiMap<String, Class<*>> = HashBiMap.create()
private val protoToServiceMap: BiMap<String, Class<*>> = HashBiMap.create()

// HashMap is faster than mutableMapOf(LinkedHashMap)
private val fileInfoMap: MutableMap<String, FileDescriptorProto> = hashMapOf()
Expand All @@ -36,6 +37,14 @@ object ProtoTypes {
protoToClassMap[protoType] = kotlinType.java
}

fun registerService(protoType: String, kotlinType: Class<*>) {
protoToServiceMap[protoType] = kotlinType
}

fun registerService(protoType: String, kotlinType: KClass<*>) {
protoToServiceMap[protoType] = kotlinType.java
}

private fun registerFileSymbol(file: FileDescriptorProto) {
val packageName = file.`package`
fileInfoMap[file.name] = file
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package com.bybutter.sisyphus.protobuf.gradle

import com.bybutter.sisyphus.io.replaceExtensionName
import com.bybutter.sisyphus.protobuf.ProtoFileMeta
import com.bybutter.sisyphus.protobuf.compiler.ProtobufGenerateContext
import com.google.protobuf.DescriptorProtos
Expand Down Expand Up @@ -63,6 +64,10 @@ open class ProtoGenerateTask : SourceTask() {
for (generateResult in result) {
generateResult.file.writeTo(output.toPath())
generateResult.implFile.writeTo(implOutput.toPath())

val descFile = Paths.get(resourceOutput.toPath().toString(), generateResult.descriptor.name.replaceExtensionName("proto", "pb"))
Files.createDirectories(descFile.parent)
Files.write(descFile, generateResult.descriptor.toProto())
}

val fileMetas = Paths.get(resourceOutput.toPath().toString(), "META-INF/services/${ProtoFileMeta::class.java.name}")
Expand Down
Original file line number Diff line number Diff line change
@@ -1,14 +1,11 @@
package com.bybutter.sisyphus.protobuf.compiler

import com.bybutter.sisyphus.api.resourceDefinition
import com.bybutter.sisyphus.data.gzip
import com.bybutter.sisyphus.data.ungzip
import com.bybutter.sisyphus.io.replaceExtensionName
import com.bybutter.sisyphus.protobuf.ProtoFileMeta
import com.bybutter.sisyphus.protobuf.ProtoTypes
import com.bybutter.sisyphus.protobuf.primitives.FileDescriptorProto
import com.bybutter.sisyphus.protobuf.primitives.FileOptions
import com.bybutter.sisyphus.security.base64
import com.bybutter.sisyphus.security.base64Decode
import com.bybutter.sisyphus.string.toPascalCase
import com.google.protobuf.DescriptorProtos
import com.squareup.kotlinpoet.ClassName
Expand All @@ -23,7 +20,7 @@ import java.nio.file.Paths
class FileGenerator(override val parent: ProtobufGenerateContext, val descriptor: DescriptorProtos.FileDescriptorProto) : ProtobufElement() {
override val protoName: String = descriptor.`package`.takeIf { it.isNotEmpty() } ?: ""
override val kotlinName: String = parent.packageMapping[protoName]
?: descriptor.options?.javaPackage?.takeIf { it.isNotEmpty() } ?: protoName
?: descriptor.options?.javaPackage?.takeIf { it.isNotEmpty() } ?: protoName

val protoFileNameWithoutExtension = Paths.get(descriptor.name).toFile().nameWithoutExtension
val kotlinFileNameWithoutExtension = protoFileNameWithoutExtension.toPascalCase()
Expand All @@ -34,6 +31,7 @@ class FileGenerator(override val parent: ProtobufGenerateContext, val descriptor
val fileMeta = "${kotlinFileNameWithoutExtension}FileMeta"
val internalKotlinName = "proto.internal.$kotlinName"
val fileMetaTypeName = ClassName(internalKotlinName, fileMeta)
val fileDescPath = protoFilePath.replaceExtensionName("proto", "pb")

override val documentation: String = """Generated by the protocol buffer dto plugin. DO NOT EDIT!
|file: $kotlinFilePath
Expand Down Expand Up @@ -69,7 +67,7 @@ class FileGenerator(override val parent: ProtobufGenerateContext, val descriptor

fun generate(): FileSpec {
val builder = FileSpec.builder(kotlinName, kotlinFileNameWithoutExtension)
.addComment(documentation)
.addComment(documentation)

for (child in children) {
when (child) {
Expand Down Expand Up @@ -100,52 +98,50 @@ class FileGenerator(override val parent: ProtobufGenerateContext, val descriptor

fun generateInternal(): FileSpec {
return FileSpec.builder(internalKotlinName, kotlinFileNameWithoutExtension)
.addType(generateFileMeta())
.apply {
for (child in children) {
when (child) {
is MessageGenerator -> {
addType(child.generateMutable())
addType(child.generateImpl())
addType(child.generateSupport())
}
is ResourceNameParentGenerator -> {
addType(child.generateSupport())
}
is ServiceGenerator -> {
addType(child.generateSupport())
.addType(generateFileMeta())
.apply {
for (child in children) {
when (child) {
is MessageGenerator -> {
addType(child.generateMutable())
addType(child.generateImpl())
addType(child.generateSupport())
}
is ResourceNameParentGenerator -> {
addType(child.generateSupport())
}
is ServiceGenerator -> {
addType(child.generateSupport())
}
}
}
}
}
.build()
.build()
}

fun generateFileMeta(): TypeSpec {
return TypeSpec.objectBuilder(fileMeta)
.addSuperinterface(ProtoFileMeta::class)
.addProperty(
PropertySpec.builder("name", String::class)
.addModifiers(KModifier.OVERRIDE)
.initializer("%S", protoFilePath)
.build()
)
.addProperty(
PropertySpec.builder("descriptor", FileDescriptorProto::class)
.addModifiers(KModifier.OVERRIDE)
.initializer("%T.parse(%S.%M().%M())",
FileDescriptorProto::class,
descriptor.toBuilder().clearSourceCodeInfo().build().toByteArray().gzip().base64(),
String::base64Decode.asMemberName(),
ByteArray::ungzip.asMemberName())
.build()
)
.addInitializerBlock(
buildCodeBlock {
registerTypes(this, this@FileGenerator)
}
)
.build()
.addSuperinterface(ProtoFileMeta::class)
.addProperty(
PropertySpec.builder("name", String::class)
.addModifiers(KModifier.OVERRIDE)
.initializer("%S", protoFilePath)
.build()
)
.addProperty(
PropertySpec.builder("descriptor", FileDescriptorProto::class)
.addModifiers(KModifier.OVERRIDE)
.initializer("%T.parse(this.javaClass.classLoader.getResourceAsStream(%S))",
FileDescriptorProto::class,
"$fileDescPath")
.build()
)
.addInitializerBlock(
buildCodeBlock {
registerTypes(this, this@FileGenerator)
}
)
.build()
}

private fun registerTypes(builder: CodeBlock.Builder, element: ProtobufElement) {
Expand All @@ -162,6 +158,9 @@ class FileGenerator(override val parent: ProtobufGenerateContext, val descriptor
is EnumGenerator -> {
builder.addStatement("%T.registerProtoType(%S, %T::class)", ProtoTypes::class, element.fullProtoName, element.kotlinType)
}
is ServiceGenerator -> {
builder.addStatement("%T.registerService(%S, %T::class)", ProtoTypes::class, element.fullProtoName, element.serviceType)
}
else -> {
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package com.bybutter.sisyphus.protobuf.compiler

import com.bybutter.sisyphus.protobuf.primitives.FileDescriptorProto
import com.google.protobuf.DescriptorProtos
import com.squareup.kotlinpoet.FileSpec

Expand All @@ -26,9 +27,9 @@ class ProtobufGenerateContext : ProtobufElement() {
return children.mapNotNull { it as? FileGenerator }.filter {
source.contains(it.protoFilePath)
}.map {
GeneratingResult(it.generate(), it.generateInternal(), it.fileMetaTypeName.toString())
GeneratingResult(it.generate(), it.generateInternal(), FileDescriptorProto.parse(it.descriptor.toByteArray()), it.fileMetaTypeName.toString())
}
}
}

data class GeneratingResult(val file: FileSpec, val implFile: FileSpec, val fileMeta: String)
data class GeneratingResult(val file: FileSpec, val implFile: FileSpec, val descriptor: FileDescriptorProto, val fileMeta: String)

0 comments on commit cb9937b

Please sign in to comment.