chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,367 @@
|
||||
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()),
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user