collectiveidea/protoc-gen-twirp_ruby

# frozen_string_literal: true

require "stringio"
require "google/protobuf/compiler/plugin_pb"
require "twirp/protoc_plugin/core_ext/string/snake_case"
require "twirp/protoc_plugin/descriptor_ext/service_descriptor_proto_ext"

module Twirp
  module ProtocPlugin
    class CodeGenerator
      # @param proto_file [Google::Protobuf::FileDescriptorProto]
      # @param relative_ruby_protobuf [String] e.g. "example_rb.pb"
      # @param options [Hash{Symbol => Symbol}]
      #   * :generate [Symbol] one of: :service, :client, or :both.
      def initialize(proto_file, relative_ruby_protobuf, options)
        @proto_file = proto_file
        @relative_ruby_protobuf = relative_ruby_protobuf
        @options = options
      end

      # @return [String] the generated Twirp::Ruby code for the proto_file
      def generate
        output = StringIO.new
        output << <<~START
          # frozen_string_literal: true
          
          # Generated by the protoc-gen-twirp_ruby gem v#{VERSION}. DO NOT EDIT!
          # source: #{@proto_file.name}

          require "twirp"
          require_relative "#{@relative_ruby_protobuf}"
      
        START

        indent_level = 0
        modules = @proto_file.ruby_module.delete_prefix("::").split("::")

        modules.each do |mod|
          output << line("module #{mod}", indent_level)
          indent_level += 1
        end

        @proto_file.service.each_with_index do |service, index| # service: <Google::Protobuf::ServiceDescriptorProto>
          # Add newline between definitions when multiple are generated
          output << "\n" if index > 0

          if %i[service both].include?(@options[:generate])
            generate_service_class(output, indent_level, service)
          end

          if @options[:generate] == :both
            # Space between generated service and client when generating both
            output << "\n"

            generate_client_class_for_service(output, indent_level, service)
          elsif @options[:generate] == :client
            # When generating only the client, we can't use the `client_for` DSL.
            generate_client_class_standalone(output, indent_level, service)
          end
        end

        modules.each do |_|
          indent_level -= 1
          output << line("end", indent_level)
        end

        output.string
      end

      private

      # Format a string by adding a trailing new line and indenting 2 spaces
      # for every indent level.
      #
      # @param input [String] the input string to format
      # @param indent_level [Integer] the number of double spaces to indent. Default 0.
      # @return [String] the input properly indented with a tailing newline added
      def line(input, indent_level = 0)
        "#{"  " * indent_level}#{input}\n"
      end

      # Generates a Twirp::Service subclass for the given service class, adding the
      # string to the output.
      #
      # @param output [#<<] the output to append the generated service code to
      # @param indent_level [Integer] the number of double spaces to indent the generated code by
      # @param service [Google::Protobuf::ServiceDescriptorProto]
      # @return [void]
      def generate_service_class(output, indent_level, service)
        output << line("class #{service.service_class_name} < ::Twirp::Service", indent_level)
        output << line("  package \"#{@proto_file.package}\"", indent_level) unless @proto_file.package.to_s.empty?
        output << line("  service \"#{service.name}\"", indent_level)
        service["method"].each do |method| # method: <Google::Protobuf::MethodDescriptorProto>
          input_type = @proto_file.ruby_type_for(method.input_type)
          output_type = @proto_file.ruby_type_for(method.output_type)
          ruby_method_name = method.name.snake_case

          output << line("  rpc :#{method.name}, #{input_type}, #{output_type}, ruby_method: :#{ruby_method_name}", indent_level)
        end
        output << line("end", indent_level)
      end

      # Generates a Twirp::Client subclass for the given service class, adding the
      # string to the output.
      #
      # @param output [#<<] the output to append the generated service code to
      # @param indent_level [Integer] the number of double spaces to indent the generated code by
      # @param service [Google::Protobuf::ServiceDescriptorProto]
      # @return [void]
      def generate_client_class_for_service(output, indent_level, service)
        output << line("class #{service.client_class_name} < ::Twirp::Client", indent_level)
        output << line("  client_for #{service.service_class_name}", indent_level)
        output << line("end", indent_level)
      end

      # Generates a Twirp::Client subclass standalone, without using the `client_for` DSL because
      # there is no corresponding service to reference.
      #
      # This essentially in-lines the `client_for` logic from
      # https://github.com/arthurnn/twirp-ruby/blob/v1.11.0/lib/twirp/client.rb#L31
      #
      # @param output [#<<] the output to append the generated service code to
      # @param indent_level [Integer] the number of double spaces to indent the generated code by
      # @param service [Google::Protobuf::ServiceDescriptorProto]
      # @return [void]
      def generate_client_class_standalone(output, indent_level, service)
        output << line("class #{service.client_class_name} < ::Twirp::Client", indent_level)
        output << line("  package \"#{@proto_file.package}\"", indent_level) unless @proto_file.package.to_s.empty?
        output << line("  service \"#{service.name}\"", indent_level)
        service["method"].each do |method| # method: <Google::Protobuf::MethodDescriptorProto>
          input_type = @proto_file.ruby_type_for(method.input_type)
          output_type = @proto_file.ruby_type_for(method.output_type)
          ruby_method_name = method.name.snake_case

          # TRICKY: The service `rpc` DSL accepts a method symbol, but the client `rpc` DSL expects a string.
          output << line("  rpc \"#{method.name}\", #{input_type}, #{output_type}, ruby_method: :#{ruby_method_name}", indent_level)
        end
        output << line("end", indent_level)
      end
    end
  end
end