wehub-resource-sync
368 lines
10 KiB
Go
368 lines
10 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"regexp"
|
|
"sort"
|
|
"strings"
|
|
|
|
log "golang.org/x/exp/slog"
|
|
"google.golang.org/protobuf/proto"
|
|
descriptor "google.golang.org/protobuf/types/descriptorpb"
|
|
"google.golang.org/protobuf/types/pluginpb"
|
|
)
|
|
|
|
const pluginVersion = "0.1.0.dev2"
|
|
|
|
func init() {
|
|
ll := log.New(log.NewTextHandler(os.Stderr, &log.HandlerOptions{
|
|
Level: log.LevelDebug,
|
|
ReplaceAttr: func(groups []string, a log.Attr) log.Attr {
|
|
if a.Key == log.TimeKey && len(groups) == 0 {
|
|
return log.Attr{}
|
|
}
|
|
return a
|
|
},
|
|
}))
|
|
log.SetDefault(ll.With(log.Int("pid", os.Getpid())))
|
|
}
|
|
|
|
func main() {
|
|
if len(os.Args) == 2 && os.Args[1] == "--version" {
|
|
fmt.Fprintln(os.Stdout, pluginVersion)
|
|
os.Exit(0)
|
|
}
|
|
|
|
f := func(plugin *Plugin) error {
|
|
for _, f := range plugin.filesToGenerate {
|
|
generate(plugin, f)
|
|
}
|
|
return nil
|
|
}
|
|
if err := run(f); err != nil {
|
|
fmt.Fprintf(os.Stderr, "%s: %v\n", filepath.Base(os.Args[0]), err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func run(f func(*Plugin) error) error {
|
|
in, err := io.ReadAll(os.Stdin)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req := &pluginpb.CodeGeneratorRequest{}
|
|
if err := proto.Unmarshal(in, req); err != nil {
|
|
return err
|
|
}
|
|
gen, err := newPlugin(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := f(gen); err != nil {
|
|
gen.Error(err)
|
|
}
|
|
resp := gen.Response()
|
|
out, err := proto.Marshal(resp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = os.Stdout.Write(out)
|
|
return err
|
|
}
|
|
|
|
func newPlugin(req *pluginpb.CodeGeneratorRequest) (*Plugin, error) {
|
|
gen := &Plugin{
|
|
request: req,
|
|
filesByPackage: make(map[string]*descriptor.FileDescriptorProto),
|
|
filesByPath: make(map[string]*descriptor.FileDescriptorProto),
|
|
messagesByType: make(map[string]*descriptor.DescriptorProto),
|
|
}
|
|
|
|
for _, f := range gen.request.ProtoFile {
|
|
name := f.GetName()
|
|
|
|
pkg := f.GetPackage()
|
|
log.Debug("ProtoFile",
|
|
log.String("name", name),
|
|
log.String("pkg", pkg),
|
|
)
|
|
// if _, ok := gen.filesByPackage[pkg]; ok {
|
|
// return nil, fmt.Errorf("duplicate package: %q", name)
|
|
// }
|
|
gen.filesByPackage[pkg] = f
|
|
if _, ok := gen.filesByPath[name]; ok {
|
|
return nil, fmt.Errorf("duplicate file name: %q", name)
|
|
}
|
|
gen.filesByPath[name] = f
|
|
for _, msg := range f.GetMessageType() {
|
|
msgKey := f.GetPackage() + "." + msg.GetName()
|
|
if _, ok := gen.messagesByType[msgKey]; ok {
|
|
return nil, fmt.Errorf("duplicate message: %q", msgKey)
|
|
}
|
|
gen.messagesByType[msgKey] = msg
|
|
log.Debug("MessageType",
|
|
log.String("name", msgKey),
|
|
)
|
|
}
|
|
gen.files = append(gen.files, f)
|
|
}
|
|
|
|
for _, name := range req.FileToGenerate {
|
|
if f, ok := gen.filesByPath[name]; ok {
|
|
log.Debug("FileToGenerate",
|
|
log.String("name", name),
|
|
log.Any("deps", f.Dependency),
|
|
log.Int("services", len(f.Service)),
|
|
)
|
|
if len(f.Service) > 0 {
|
|
gen.filesToGenerate = append(gen.filesToGenerate, f)
|
|
}
|
|
} else {
|
|
return nil, fmt.Errorf("missing file: %q", name)
|
|
}
|
|
}
|
|
|
|
return gen, nil
|
|
}
|
|
|
|
type Plugin struct {
|
|
request *pluginpb.CodeGeneratorRequest
|
|
|
|
files []*descriptor.FileDescriptorProto
|
|
filesByPackage map[string]*descriptor.FileDescriptorProto
|
|
filesByPath map[string]*descriptor.FileDescriptorProto
|
|
messagesByType map[string]*descriptor.DescriptorProto
|
|
filesToGenerate []*descriptor.FileDescriptorProto
|
|
|
|
generatedFiles []*pluginpb.CodeGeneratorResponse_File
|
|
|
|
err error
|
|
}
|
|
|
|
func (gen *Plugin) Response() *pluginpb.CodeGeneratorResponse {
|
|
resp := &pluginpb.CodeGeneratorResponse{}
|
|
if gen.err != nil {
|
|
resp.Error = Ptr(gen.err.Error())
|
|
return resp
|
|
}
|
|
resp.File = gen.generatedFiles
|
|
return resp
|
|
}
|
|
|
|
func (gen *Plugin) Error(err error) {
|
|
if gen.err == nil {
|
|
gen.err = err
|
|
}
|
|
}
|
|
|
|
func Ptr[T any](v T) *T {
|
|
return &v
|
|
}
|
|
|
|
func print(buf *strings.Builder, tpl string, args ...interface{}) {
|
|
buf.WriteString(fmt.Sprintf(tpl, args...))
|
|
buf.WriteByte('\n')
|
|
}
|
|
|
|
func getPackage(path string) string {
|
|
return strings.ReplaceAll(filepath.Dir(path), "/", ".")
|
|
}
|
|
|
|
func getModule(path string) string {
|
|
path = filepath.Base(path)
|
|
ext := filepath.Ext(path)
|
|
return strings.TrimSuffix(path, ext)
|
|
}
|
|
|
|
func getProtoModule(path string) string {
|
|
return getModule(path) + "_pb2"
|
|
}
|
|
|
|
func getConnectModule(path string) string {
|
|
return getModule(path) + "_connect"
|
|
}
|
|
|
|
func getProtoModuleAlias(path string) string {
|
|
path = getPackage(path) + "." + getProtoModule(path)
|
|
path = strings.ReplaceAll(path, "_", "__")
|
|
path = strings.ReplaceAll(path, ".", "_dot_")
|
|
return path
|
|
}
|
|
|
|
func getServiceName(svc *descriptor.ServiceDescriptorProto) string {
|
|
return svc.GetName() + "Name"
|
|
}
|
|
|
|
func getServiceClient(svc *descriptor.ServiceDescriptorProto) string {
|
|
return svc.GetName() + "Client"
|
|
}
|
|
|
|
func getServiceBasePath(file *descriptor.FileDescriptorProto, svc *descriptor.ServiceDescriptorProto) string {
|
|
return file.GetPackage() + "." + svc.GetName()
|
|
}
|
|
|
|
func getMethodProperty(m *descriptor.MethodDescriptorProto) string {
|
|
return "_" + toSnakeCase(m.GetName())
|
|
}
|
|
|
|
func getMethodType(m *descriptor.MethodDescriptorProto) string {
|
|
switch {
|
|
case m.GetClientStreaming() && m.GetServerStreaming():
|
|
return "bidi_stream"
|
|
case m.GetClientStreaming():
|
|
return "client_stream"
|
|
case m.GetServerStreaming():
|
|
return "server_stream"
|
|
default:
|
|
return "unary"
|
|
}
|
|
}
|
|
|
|
func splitPackageType(path string) (string, string) {
|
|
lastDot := strings.LastIndexByte(path, '.')
|
|
return path[:lastDot], path[lastDot+1:]
|
|
}
|
|
|
|
func resolveMessageFromMethod(gen *Plugin, m *descriptor.MethodDescriptorProto) (string, string) {
|
|
fullyQualifiedName := m.GetOutputType()[1:] // strip prefixed "."
|
|
pkgName, msgName := splitPackageType(fullyQualifiedName)
|
|
filename := gen.filesByPackage[pkgName].GetName()
|
|
return filename, msgName
|
|
}
|
|
|
|
func resolveInputFromMethod(gen *Plugin, m *descriptor.MethodDescriptorProto) (string, string) {
|
|
fullyQualifiedName := m.GetInputType()[1:] // strip prefixed "."
|
|
pkgName, msgName := splitPackageType(fullyQualifiedName)
|
|
filename := gen.filesByPackage[pkgName].GetName()
|
|
return filename, msgName
|
|
}
|
|
|
|
func getResponseType(gen *Plugin, m *descriptor.MethodDescriptorProto) string {
|
|
filename, msgName := resolveMessageFromMethod(gen, m)
|
|
|
|
log.Debug("ResponseType",
|
|
log.String("msg", msgName),
|
|
log.String("import", filename),
|
|
log.String("alias", getProtoModuleAlias(filename)),
|
|
)
|
|
return getProtoModuleAlias(filename) + "." + msgName
|
|
}
|
|
|
|
func getRequestType(gen *Plugin, m *descriptor.MethodDescriptorProto) string {
|
|
filename, msgName := resolveInputFromMethod(gen, m)
|
|
|
|
log.Debug("RequestType",
|
|
log.String("msg", msgName),
|
|
log.String("import", filename),
|
|
log.String("alias", getProtoModuleAlias(filename)),
|
|
)
|
|
return getProtoModuleAlias(filename) + "." + msgName
|
|
}
|
|
|
|
var (
|
|
matchFirstCap = regexp.MustCompile("(.)([A-Z][a-z]+)")
|
|
matchAllCap = regexp.MustCompile("([a-z0-9])([A-Z])")
|
|
)
|
|
|
|
func toSnakeCase(str string) string {
|
|
snake := matchFirstCap.ReplaceAllString(str, "${1}_${2}")
|
|
snake = matchAllCap.ReplaceAllString(snake, "${1}_${2}")
|
|
return strings.ToLower(snake)
|
|
}
|
|
|
|
func generate(gen *Plugin, file *descriptor.FileDescriptorProto) {
|
|
filename := file.GetName()
|
|
|
|
dir := filepath.Dir(filename)
|
|
pkgName := getPackage(filename)
|
|
modName := getModule(filename)
|
|
|
|
log.Debug("Generate",
|
|
log.String("name", filename),
|
|
log.String("pkg", pkgName),
|
|
log.String("mod", modName),
|
|
)
|
|
|
|
b := new(strings.Builder)
|
|
|
|
depsUniq := make(map[string]struct{})
|
|
for _, svc := range file.Service {
|
|
for _, method := range svc.Method {
|
|
filename, _ := resolveMessageFromMethod(gen, method)
|
|
depsUniq[filename] = struct{}{}
|
|
}
|
|
}
|
|
|
|
deps := make([]string, 0, len(depsUniq))
|
|
for dep := range depsUniq {
|
|
deps = append(deps, dep)
|
|
}
|
|
sort.Strings(deps)
|
|
|
|
print(b, "# Code generated by protoc-gen-connect-python %s, DO NOT EDIT.", pluginVersion)
|
|
|
|
print(b, "from typing import Any, Generator, Coroutine, AsyncGenerator, Optional")
|
|
print(b, "from httpcore import ConnectionPool, AsyncConnectionPool")
|
|
print(b, "")
|
|
|
|
print(b, "import e2b_connect as connect")
|
|
if len(deps) > 0 {
|
|
print(b, "")
|
|
for _, dep := range deps {
|
|
print(b, "from %s import %s as %s", getPackage(dep), getProtoModule(dep), getProtoModuleAlias(dep))
|
|
}
|
|
}
|
|
print(b, "")
|
|
|
|
for _, svc := range file.Service {
|
|
print(b, `%s = "%s"`, getServiceName(svc), getServiceBasePath(file, svc))
|
|
}
|
|
|
|
for _, svc := range file.Service {
|
|
print(b, "")
|
|
print(b, "")
|
|
print(b, `class %s:`, getServiceClient(svc))
|
|
print(b, " def __init__(self, base_url: str, *, pool: Optional[ConnectionPool] = None, async_pool: Optional[AsyncConnectionPool] = None, compressor=None, json=False, **opts):")
|
|
if len(svc.Method) == 0 {
|
|
print(b, " pass")
|
|
continue
|
|
}
|
|
for _, method := range svc.Method {
|
|
print(b, " self.%s = connect.Client(", getMethodProperty(method))
|
|
print(b, " pool=pool,")
|
|
print(b, " async_pool=async_pool,")
|
|
print(b, ` url=f"{base_url}/{%s}/%s",`, getServiceName(svc), method.GetName())
|
|
print(b, ` response_type=%s,`, getResponseType(gen, method))
|
|
print(b, ` compressor=compressor,`)
|
|
print(b, ` json=json,`)
|
|
print(b, ` **opts`)
|
|
print(b, " )")
|
|
}
|
|
for _, method := range svc.Method {
|
|
print(b, "")
|
|
|
|
if method.GetServerStreaming() {
|
|
print(b, " def %s(self, req: %s , **opts) -> Generator[%s, Any, None]:", toSnakeCase(method.GetName()), getRequestType(gen, method), getResponseType(gen, method))
|
|
print(b, " return self.%s.call_%s(req, **opts)", getMethodProperty(method), getMethodType(method))
|
|
print(b, "")
|
|
print(b, " def a%s(self, req: %s , **opts) -> AsyncGenerator[%s, Any]:", toSnakeCase(method.GetName()), getRequestType(gen, method), getResponseType(gen, method))
|
|
print(b, " return self.%s.acall_%s(req, **opts)", getMethodProperty(method), getMethodType(method))
|
|
} else {
|
|
print(b, " def %s(self, req: %s, **opts) -> %s:", toSnakeCase(method.GetName()), getRequestType(gen, method), getResponseType(gen, method))
|
|
print(b, " return self.%s.call_%s(req, **opts)", getMethodProperty(method), getMethodType(method))
|
|
print(b, "")
|
|
print(b, " def a%s(self, req: %s, **opts) -> Coroutine[Any, Any, %s]:", toSnakeCase(method.GetName()), getRequestType(gen, method), getResponseType(gen, method))
|
|
print(b, " return self.%s.acall_%s(req, **opts)", getMethodProperty(method), getMethodType(method))
|
|
}
|
|
}
|
|
}
|
|
|
|
gen.generatedFiles = append(gen.generatedFiles, &pluginpb.CodeGeneratorResponse_File{
|
|
Name: Ptr(filepath.Join(dir, getConnectModule(filename)+".py")),
|
|
Content: Ptr(b.String()),
|
|
})
|
|
}
|