twirp-ruby/protoc-gen-twirp_ruby/main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"unicode"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

const Version = "v5.2.0"

func main() {
	versionFlag := flag.Bool("version", false, "print version and exit")
	flag.Parse()
	if *versionFlag {
		fmt.Println(Version)
		os.Exit(0)
	}

	g := newGenerator()
	Main(g)
}

type Generator interface {
	Generate(in *plugin.CodeGeneratorRequest) *plugin.CodeGeneratorResponse
}

func Main(g Generator) {
	req := readGenRequest(os.Stdin)
	resp := g.Generate(req)
	writeResponse(os.Stdout, resp)
}

type generator struct {
	output *bytes.Buffer
}

func newGenerator() *generator {
	return &generator{output: new(bytes.Buffer)}
}

func (g *generator) Generate(in *plugin.CodeGeneratorRequest) *plugin.CodeGeneratorResponse {

	resp := new(plugin.CodeGeneratorResponse)
	for _, name := range in.FileToGenerate {
		for _, f := range in.ProtoFile {
			if f.GetName() == name {
				respFile := g.generateFile(f)
				if respFile != nil {
					resp.File = append(resp.File, respFile)
				}
				continue
			}
		}
	}

	return resp
}

func (g *generator) generateFile(file *descriptor.FileDescriptorProto) *plugin.CodeGeneratorResponse_File {
	indent := ""
	pkgName := pkgName(file)
	g.P(`# Code generated by protoc-gen-twirp_ruby, DO NOT EDIT.`)
	g.P(`require 'twirp'`)
	g.P(``)
	if pkgName != "" {
		g.P(fmt.Sprintf("module %s", CamelCase(pkgName)))
		indent = indent + "  "
	}
	for i, service := range file.Service {
		serviceName := serviceName(service)
		g.P(fmt.Sprintf("%sclass %sService < Twirp::Service", indent, CamelCase(serviceName)))
		if pkgName != "" {
			g.P(fmt.Sprintf(`%s  package "%s"`, indent, pkgName))
		}
		g.P(fmt.Sprintf(`%s  service "%s"`, indent, serviceName))
		for _, method := range service.GetMethod() {
			methName := methodName(method)
			inputName := methodInputName(method)
			outputName := methodOutputName(method)
			g.P(fmt.Sprintf(`%s  rpc "%s", %s, %s, handler_method: "%s"`,
				indent, methName, inputName, outputName, SnakeCase(methName)))
		}
		g.P(fmt.Sprintf(`%send`, indent))
		if i < len(file.Service)-1 {
			g.P(``)
		}
	}
	if pkgName != "" {
		g.P(`end`)
	}

	resp := new(plugin.CodeGeneratorResponse_File)
	resp.Name = proto.String(rubyFileName(file))
	resp.Content = proto.String(g.output.String())
	g.output.Reset()

	return resp
}

func (g *generator) P(args ...string) {
	for _, v := range args {
		g.output.WriteString(v)
	}
	g.output.WriteByte('\n')
}

func rubyFileName(f *descriptor.FileDescriptorProto) string {
	name := *f.Name
	if ext := path.Ext(name); ext == ".proto" || ext == ".protodevel" {
		name = name[:len(name)-len(ext)]
	}
	name += "_twirp.rb"
	return name
}

func pkgName(file *descriptor.FileDescriptorProto) string {
	return file.GetPackage()
}

func serviceName(service *descriptor.ServiceDescriptorProto) string {
	return service.GetName()
}

func methodName(method *descriptor.MethodDescriptorProto) string {
	return method.GetName()
}

// methodInputName returns the basename of the input type of a method in snake
// case.
func methodInputName(meth *descriptor.MethodDescriptorProto) string {
	fullName := meth.GetInputType()
	split := strings.Split(fullName, ".")
	return split[len(split)-1]
}

// methodInputName returns the basename of the input type of a method in snake
// case.
func methodOutputName(meth *descriptor.MethodDescriptorProto) string {
	fullName := meth.GetOutputType()
	split := strings.Split(fullName, ".")
	return split[len(split)-1]
}

func Fail(msgs ...string) {
	s := strings.Join(msgs, " ")
	log.Print("error:", s)
	os.Exit(1)
}

// SnakeCase converts a string from CamelCase to snake_case.
func SnakeCase(s string) string {
	var buf bytes.Buffer
	for i, r := range s {
		if unicode.IsUpper(r) && i > 0 {
			fmt.Fprintf(&buf, "_")
		}
		r = unicode.ToLower(r)
		fmt.Fprintf(&buf, "%c", r)
	}
	return buf.String()
}

func readGenRequest(r io.Reader) *plugin.CodeGeneratorRequest {
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		Fail(err.Error(), "reading input")
	}

	req := new(plugin.CodeGeneratorRequest)
	if err = proto.Unmarshal(data, req); err != nil {
		Fail(err.Error(), "parsing input proto")
	}

	if len(req.FileToGenerate) == 0 {
		Fail("no files to generate")
	}

	return req
}

func writeResponse(w io.Writer, resp *plugin.CodeGeneratorResponse) {
	data, err := proto.Marshal(resp)
	if err != nil {
		Fail(err.Error(), "marshaling response")
	}
	_, err = w.Write(data)
	if err != nil {
		Fail(err.Error(), "writing response")
	}
}

// CamelCase converts a string from snake_case to CamelCased.
//
// If there is an interior underscore followed by a lower case letter, drop the
// underscore and convert the letter to upper case. There is a remote
// possibility of this rewrite causing a name collision, but it's so remote
// we're prepared to pretend it's nonexistent - since the C++ generator
// lowercases names, it's extremely unlikely to have two fields with different
// capitalizations. In short, _my_field_name_2 becomes XMyFieldName_2.
func CamelCase(s string) string {
	if s == "" {
		return ""
	}
	t := make([]byte, 0, 32)
	i := 0
	if s[0] == '_' {
		// Need a capital letter; drop the '_'.
		t = append(t, 'X')
		i++
	}
	// Invariant: if the next letter is lower case, it must be converted
	// to upper case.
	//
	// That is, we process a word at a time, where words are marked by _ or upper
	// case letter. Digits are treated as words.
	for ; i < len(s); i++ {
		c := s[i]
		if c == '_' && i+1 < len(s) && isASCIILower(s[i+1]) {
			continue // Skip the underscore in s.
		}
		if isASCIIDigit(c) {
			t = append(t, c)
			continue
		}
		// Assume we have a letter now - if not, it's a bogus identifier. The next
		// word is a sequence of characters that must start upper case.
		if isASCIILower(c) {
			c ^= ' ' // Make it a capital letter.
		}
		t = append(t, c) // Guaranteed not lower case.
		// Accept lower case sequence that follows.
		for i+1 < len(s) && isASCIILower(s[i+1]) {
			i++
			t = append(t, s[i])
		}
	}
	return string(t)
}

// Is c an ASCII lower-case letter?
func isASCIILower(c byte) bool {
	return 'a' <= c && c <= 'z'
}

// Is c an ASCII digit?
func isASCIIDigit(c byte) bool {
	return '0' <= c && c <= '9'
}