[ProtobufCodeGen] Use the generated input and output type names (#1792)

stefanadranca · glbrntt · web-flow · commit daa008e437c8 · 2024-02-06T15:07:44.000Z
Motivation:

protoc-gen-swift generates message type names that are different from the ones in the file descriptors. The CodeGenerationRequest object should contain the generated names for the ProtobufCodeGen.

Modifications:

- the parser uses a SwiftProtobufNamer to get the generated type names
and adds them in the CodeGenerationRequest objects
- propagated the change and modified tests accordingly

Result:

The input and output type names will be the generated ones and the generated code will compile.


---------

Co-authored-by: George Barnett &lt;gbarnett@apple.com&gt;
diff --git a/Sources/GRPCProtobufCodeGen/ProtobufCodeGenParser.swift b/Sources/GRPCProtobufCodeGen/ProtobufCodeGenParser.swift
@@ -22,9 +22,19 @@ import struct GRPCCodeGen.CodeGenerationRequest
 
 /// Parses a ``FileDescriptor`` object into a ``CodeGenerationRequest`` object.
 internal struct ProtobufCodeGenParser {
-  internal init() {}
-  internal func parse(input: FileDescriptor) throws -> CodeGenerationRequest {
-    var header = input.header
+  let input: FileDescriptor
+  let namer: SwiftProtobufNamer
+
+  internal init(input: FileDescriptor, protoFileModuleMappings: ProtoFileToModuleMappings) {
+    self.input = input
+    self.namer = SwiftProtobufNamer(
+      currentFile: input,
+      protoFileToModuleMappings: protoFileModuleMappings
+    )
+  }
+
+  internal func parse() throws -> CodeGenerationRequest {
+    var header = self.input.header
     // Ensuring there is a blank line after the header.
     if !header.isEmpty && !header.hasSuffix("\n\n") {
       header.append("\n")
@@ -34,13 +44,13 @@ internal struct ProtobufCodeGenParser {
       // swift-format-ignore-file
       //
       // Generated by the gRPC Swift generator plugin for the protocol buffer compiler.
-      // Source: \(input.name)
+      // Source: \(self.input.name)
       //
       // For information on using the generated types, please see the documentation:
       //   https://github.com/grpc/grpc-swift
 
       """
-    var dependencies = input.dependencies.map {
+    var dependencies = self.input.dependencies.map {
       CodeGenerationRequest.Dependency(module: $0.name)
     }
     dependencies.append(CodeGenerationRequest.Dependency(module: "GRPCProtobuf"))
@@ -50,12 +60,16 @@ internal struct ProtobufCodeGenParser {
     let lookupDeserializer: (String) -> String = { messageType in
       "ProtobufDeserializer<\(messageType)>()"
     }
-    let services = input.services.map {
-      CodeGenerationRequest.ServiceDescriptor(descriptor: $0, package: input.package)
+    let services = self.input.services.map {
+      CodeGenerationRequest.ServiceDescriptor(
+        descriptor: $0,
+        package: input.package,
+        protobufNamer: self.namer
+      )
     }
 
     return CodeGenerationRequest(
-      fileName: input.name,
+      fileName: self.input.name,
       leadingTrivia: header + leadingTrivia,
       dependencies: dependencies,
       services: services,
@@ -66,9 +80,16 @@ internal struct ProtobufCodeGenParser {
 }
 
 extension CodeGenerationRequest.ServiceDescriptor {
-  fileprivate init(descriptor: ServiceDescriptor, package: String) {
+  fileprivate init(
+    descriptor: ServiceDescriptor,
+    package: String,
+    protobufNamer: SwiftProtobufNamer
+  ) {
     let methods = descriptor.methods.map {
-      CodeGenerationRequest.ServiceDescriptor.MethodDescriptor(descriptor: $0)
+      CodeGenerationRequest.ServiceDescriptor.MethodDescriptor(
+        descriptor: $0,
+        protobufNamer: protobufNamer
+      )
     }
     let name = CodeGenerationRequest.Name(
       base: descriptor.name,
@@ -86,7 +107,7 @@ extension CodeGenerationRequest.ServiceDescriptor {
 }
 
 extension CodeGenerationRequest.ServiceDescriptor.MethodDescriptor {
-  fileprivate init(descriptor: MethodDescriptor) {
+  fileprivate init(descriptor: MethodDescriptor, protobufNamer: SwiftProtobufNamer) {
     let name = CodeGenerationRequest.Name(
       base: descriptor.name,
       generatedUpperCase: NamingUtils.toUpperCamelCase(descriptor.name),
@@ -98,8 +119,8 @@ extension CodeGenerationRequest.ServiceDescriptor.MethodDescriptor {
       name: name,
       isInputStreaming: descriptor.clientStreaming,
       isOutputStreaming: descriptor.serverStreaming,
-      inputType: descriptor.inputType.name,
-      outputType: descriptor.outputType.name
+      inputType: protobufNamer.fullName(message: descriptor.inputType),
+      outputType: protobufNamer.fullName(message: descriptor.outputType)
     )
   }
 }
diff --git a/Sources/GRPCProtobufCodeGen/ProtobufCodeGenerator.swift b/Sources/GRPCProtobufCodeGen/ProtobufCodeGenerator.swift
@@ -20,15 +20,23 @@ import SwiftProtobufPluginLibrary
 public struct ProtobufCodeGenerator {
   internal var configuration: SourceGenerator.Configuration
 
-  public init(configuration: SourceGenerator.Configuration) {
+  public init(
+    configuration: SourceGenerator.Configuration
+  ) {
     self.configuration = configuration
   }
 
-  public func generateCode(from fileDescriptor: FileDescriptor) throws -> String {
-    let parser = ProtobufCodeGenParser()
+  public func generateCode(
+    from fileDescriptor: FileDescriptor,
+    protoFileModuleMappings: ProtoFileToModuleMappings
+  ) throws -> String {
+    let parser = ProtobufCodeGenParser(
+      input: fileDescriptor,
+      protoFileModuleMappings: protoFileModuleMappings
+    )
     let sourceGenerator = SourceGenerator(configuration: self.configuration)
 
-    let codeGenerationRequest = try parser.parse(input: fileDescriptor)
+    let codeGenerationRequest = try parser.parse()
     let sourceFile = try sourceGenerator.generate(codeGenerationRequest)
     return sourceFile.contents
   }
diff --git a/Sources/protoc-gen-grpc-swift/main.swift b/Sources/protoc-gen-grpc-swift/main.swift
@@ -172,7 +172,10 @@ func main(args: [String]) throws {
           let grpcGenerator = ProtobufCodeGenerator(
             configuration: SourceGenerator.Configuration(options: options)
           )
-          grpcFile.content = try grpcGenerator.generateCode(from: fileDescriptor)
+          grpcFile.content = try grpcGenerator.generateCode(
+            from: fileDescriptor,
+            protoFileModuleMappings: options.protoToModuleMappings
+          )
         } else {
           let grpcGenerator = Generator(fileDescriptor, options: options)
           grpcFile.content = grpcGenerator.code
diff --git a/Tests/GRPCProtobufCodeGenTests/ProtobufCodeGenParserTests.swift b/Tests/GRPCProtobufCodeGenTests/ProtobufCodeGenParserTests.swift
@@ -31,9 +31,7 @@ final class ProtobufCodeGenParserTests: XCTestCase {
         """
       )
     }
-    let parsedCodeGenRequest = try ProtobufCodeGenParser().parse(
-      input: fileDescriptor
-    )
+    let parsedCodeGenRequest = try ProtobufCodeGenParser(input: fileDescriptor, protoFileModuleMappings: ProtoFileToModuleMappings()).parse()
     XCTAssertEqual(parsedCodeGenRequest.fileName, "helloworld.proto")
     XCTAssertEqual(
       parsedCodeGenRequest.leadingTrivia,
@@ -75,8 +73,8 @@ final class ProtobufCodeGenParserTests: XCTestCase {
       ),
       isInputStreaming: false,
       isOutputStreaming: false,
-      inputType: "HelloRequest",
-      outputType: "HelloReply"
+      inputType: "Helloworld_HelloRequest",
+      outputType: "Helloworld_HelloReply"
     )
     guard let method = parsedCodeGenRequest.services.first?.methods.first else { return XCTFail() }
     XCTAssertEqual(method, expectedMethod)
diff --git a/Tests/GRPCProtobufCodeGenTests/ProtobufCodeGeneratorTests.swift b/Tests/GRPCProtobufCodeGenTests/ProtobufCodeGeneratorTests.swift
@@ -60,8 +60,8 @@ final class ProtobufCodeGeneratorTests: XCTestCase {
             internal enum Greeter {
                 internal enum Method {
                     internal enum SayHello {
-                        internal typealias Input = HelloRequest
-                        internal typealias Output = HelloReply
+                        internal typealias Input = Helloworld_HelloRequest
+                        internal typealias Output = Helloworld_HelloReply
                         internal static let descriptor = MethodDescriptor(
                             service: "helloworld.Greeter",
                             method: "SayHello"
@@ -165,8 +165,8 @@ final class ProtobufCodeGeneratorTests: XCTestCase {
           public enum Greeter {
             public enum Method {
               public enum SayHello {
-                public typealias Input = HelloRequest
-                public typealias Output = HelloReply
+                public typealias Input = Helloworld_HelloRequest
+                public typealias Output = Helloworld_HelloReply
                 public static let descriptor = MethodDescriptor(
                   service: "helloworld.Greeter",
                   method: "SayHello"
@@ -254,8 +254,8 @@ final class ProtobufCodeGeneratorTests: XCTestCase {
           package enum Greeter {
             package enum Method {
               package enum SayHello {
-                package typealias Input = HelloRequest
-                package typealias Output = HelloReply
+                package typealias Input = Helloworld_HelloRequest
+                package typealias Output = Helloworld_HelloReply
                 package static let descriptor = MethodDescriptor(
                   service: "helloworld.Greeter",
                   method: "SayHello"
@@ -384,7 +384,13 @@ final class ProtobufCodeGeneratorTests: XCTestCase {
       )
     }
     let generator = ProtobufCodeGenerator(configuration: configs)
-    try XCTAssertEqualWithDiff(try generator.generateCode(from: fileDescriptor), expectedCode)
+    try XCTAssertEqualWithDiff(
+      try generator.generateCode(
+        from: fileDescriptor,
+        protoFileModuleMappings: ProtoFileToModuleMappings()
+      ),
+      expectedCode
+    )
   }
 }
 
