main.go

package main

import (
	"flag"
	"fmt"
	"io"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/pluginpb"
)

var (
	Version = "(unknown)"

	SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
	matchFirstCap     = regexp.MustCompile("([a-z0-9])([A-Z])")
	matchAllCap       = regexp.MustCompile("([A-Z])([A-Z][a-z])")

	tmpl *template.Template
)

func init() {
	tmpl = template.Must(template.New("pydantic").Parse(modelTemplate))
}

func main() {
	var flags flag.FlagSet
	_ = flags.String("root", ".", "")
	preservingProtoFieldName := flags.Bool("preserving_proto_field_name", false, "")
	autoTrimEnumPrefix := flags.Bool("auto_trim_enum_prefix", true, "")
	useIntegersForEnums := flags.Bool("use_integers_for_enums", false, "")
	disableFieldDescription := flags.Bool("disable_field_description", false, "")
	useNoneUnionSyntaxInstaedOfOptional := flags.Bool("use_none_union_syntax_instead_of_optional", false, "")

	opts := protogen.Options{
		ParamFunc: flags.Set,
	}
	opts.Run(func(gen *protogen.Plugin) error {
		protocVersion := "(unknown)"
		if v := gen.Request.GetCompilerVersion(); v != nil {
			protocVersion = fmt.Sprintf("v%v.%v.%v", v.GetMajor(), v.GetMinor(), v.GetPatch())
			if s := v.GetSuffix(); s != "" {
				protocVersion += "-" + s
			}
		}
		gen.SupportedFeatures = SupportedFeatures

		e := NewGenerator(GeneratorConfig{
			PreservingProtoFieldName:            *preservingProtoFieldName,
			AutoTrimEnumPrefix:                  *autoTrimEnumPrefix,
			UseIntegersForEnums:                 *useIntegersForEnums,
			DisableFieldDescription:             *disableFieldDescription,
			UseNoneUnionSyntaxInstaedOfOptional: *useNoneUnionSyntaxInstaedOfOptional,
		})

		for _, f := range gen.Files {
			if !f.Generate {
				continue
			}

			if err := e.processFile(f.Desc, f.Proto); err != nil {
				return err
			}
		}

		for _, f := range gen.Files {
			if !f.Generate {
				continue
			}
			filename := f.GeneratedFilenamePrefix + "_pydantic.py"
			g := gen.NewGeneratedFile(filename, f.GoImportPath)
			err := e.Generate(g)
			if err != nil {
				return fmt.Errorf("failed to write to %s: %w", filename, err)
			}
		}

		return nil
	})
}

// Template for Pydantic model
const modelTemplate = `# DO NOT EDIT. Generated by protoc-gen-pydantic.
"""
{{- $config := .Config -}}
{{- range .File.LeadingComments }}
{{ . }}
{{- end }}
"""

from enum import Enum as _Enum
from typing import Optional as _Optional

from pydantic import BaseModel as _BaseModel, Field as _Field

{{- range .ExternalImports }}

{{ . }}{{ end }}



{{- range .Enums }}

class {{ .Name }}({{ if $config.UseIntegersForEnums }}int{{ else }}str{{ end }}, _Enum):
    """
    {{- range .LeadingComments }}
    {{ . }}
    {{- end }}

    {{- range .Values }}
    {{ range .LeadingComments }}
    # {{ . }}
    {{- end }}
    {{- if $config.UseIntegersForEnums }}
    {{ .Name }} = {{ .Number }}  # {{ .Name }}
    {{ else }}
    {{ .Name }} = "{{ .Name }}"  # {{ .Number }}
    {{- end }}
    {{- end }}
    """
    {{ range .TrailingComments }}
    # {{ . }}
    {{- end }}
    {{- range .Values }}
    {{ range .LeadingComments }}
    # {{ . }}
    {{- end }}
    {{- if $config.UseIntegersForEnums }}
    {{ .Name }} = {{ .Number }}  # {{ .Name }}
    {{ else }}
    {{ .Name }} = "{{ .Name }}"  # {{ .Number }}
    {{- end }}
    {{- range .TrailingComments }}
    # {{ . }}
    {{- end }}
    {{- end }}

{{- end }}

{{- range .Messages }}

class {{ .Name }}(_BaseModel):
    """
    {{- range .LeadingComments }}
    {{ . }}
    {{- end }}

    Attributes:
    {{- range .Fields }}
      {{ .Name }} ({{ .Type }}):
    {{- range .LeadingComments }}
        {{ . }}
    {{- end }}
    {{- end }}
    """
    {{ range .TrailingComments }}
    # {{ . }}
    {{- end }}
    {{- range .Fields }}
    {{ range .LeadingComments }}
    # {{ . }}
    {{- end }}
    {{ .Name }}: "{{ .Type }}" = _Field({{ if or .Optional (ne .OneOf nil) }}None{{ else }}...{{ end }}{{ if and (not $config.DisableFieldDescription) (or (ne (len .LeadingComments) 0) (ne .OneOf nil)) }}, description="""{{- range .LeadingComments -}}
{{ . }}
{{ end }}{{ if ne .OneOf nil }}
Only one of the fields can be specified with: {{ .OneOf.FieldNames }} (oneof {{ .OneOf.Name }}){{ end }}"""{{ end }})
    {{- range .TrailingComments }}
    # {{ . }}
    {{- end }}
    {{- end }}
    {{- if eq (len .Fields) 0 }}
    pass
    {{ end }}

{{- end }}
{{- range .File.TrailingComments }}
# {{ . }}
{{- end }}
`

type EnumValue struct {
	Name             string
	Number           int32
	LeadingComments  []string
	TrailingComments []string
}

type Enum struct {
	Name             string
	Values           []EnumValue
	LeadingComments  []string
	TrailingComments []string
}

type Field struct {
	Name             string
	Type             string
	Optional         bool
	OneOf            *OneOf
	LeadingComments  []string
	TrailingComments []string
}

type OneOf struct {
	Name       string
	FieldNames []string
}

type Message struct {
	Name             string
	Fields           []Field
	LeadingComments  []string
	TrailingComments []string
}

func (m Message) TopoKey() string {
	return m.Name
}

type File struct {
	LeadingComments  []string
	TrailingComments []string
}

type generator struct {
	file            File
	enums           []Enum
	messages        []Message
	externalImports []string

	config GeneratorConfig
}

type GeneratorConfig struct {
	PreservingProtoFieldName            bool
	AutoTrimEnumPrefix                  bool
	UseIntegersForEnums                 bool
	DisableFieldDescription             bool
	UseNoneUnionSyntaxInstaedOfOptional bool
}

func NewGenerator(c GeneratorConfig) *generator {
	return &generator{
		// deps: make(map[string][]string),
		config: c,
	}
}

func (e *generator) Generate(w io.Writer) error {
	// msgs, err := topologicalSort(e.Messages, e.deps)
	// if err != nil {
	// 	return err
	// }
	err := tmpl.Execute(w, struct {
		File            File
		Enums           []Enum
		Messages        []Message
		ExternalImports []string
		Config          GeneratorConfig
	}{
		e.file,
		e.enums,
		// msgs,
		e.messages,
		e.externalImports,
		e.config,
	})
	if err != nil {
		return err
	}
	return nil
}

func (e *generator) processFile(file protoreflect.FileDescriptor, fdp *descriptorpb.FileDescriptorProto) error {
	def := File{}
	sourceCodeInfo := fdp.GetSourceCodeInfo()
	// for _, location := range sourceCodeInfo.Location {
	// 	fmt.Fprintf(os.Stderr, "location: %v => %s, %s, %s\n", location.Path, location.GetLeadingComments(), location.GetTrailingComments(), location.GetLeadingDetachedComments())
	// }
	path := []int32{12}
	def.LeadingComments, def.TrailingComments = extractComments(sourceCodeInfo, path)
	e.file = def
	for i := range file.Enums().Len() {
		ed := file.Enums().Get(i)
		// ep := fd.GetEnumType()[i]
		path := []int32{5, int32(i)}
		if err := e.processEnum(ed, sourceCodeInfo, path); err != nil {
			return err
		}
	}
	for i := range file.Messages().Len() {
		msgd := file.Messages().Get(i)
		// msgp := fd.GetMessageType()[i]
		path := []int32{4, int32(i)}
		if err := e.processMessage(msgd, sourceCodeInfo, path); err != nil {
			return err
		}
	}
	return nil
}

func (e *generator) processEnum(
	enum protoreflect.EnumDescriptor,
	sourceCodeInfo *descriptorpb.SourceCodeInfo,
	path []int32,
) error {
	def := Enum{
		Name:   resolveName(enum),
		Values: []EnumValue{},
	}
	def.LeadingComments, def.TrailingComments = extractComments(sourceCodeInfo, path)

	prefix := camelToSnakeCase(string(enum.Name())) + "_"
	for i := range enum.Values().Len() {
		v := enum.Values().Get(i)
		valueName := string(v.Name())
		if e.config.AutoTrimEnumPrefix {
			valueName = strings.TrimPrefix(valueName, prefix)
		}
		fieldpath := append(path, 2, int32(i))
		leadingComments, trailingComments := extractComments(sourceCodeInfo, fieldpath)
		def.Values = append(def.Values, EnumValue{
			Name:             valueName,
			Number:           int32(v.Number()),
			LeadingComments:  leadingComments,
			TrailingComments: trailingComments,
		})
	}

	e.enums = append(e.enums, def)
	return nil
}

func (e *generator) processMessage(
	msg protoreflect.MessageDescriptor,
	sourceCodeInfo *descriptorpb.SourceCodeInfo,
	path []int32,
) error {
	if msg.IsMapEntry() {
		return nil
	}

	// NOTE: Process nested enums and messages before the fields.
	for i, nest := range iter(msg.Enums()) {
		nest_path := append(path, 4, int32(i))
		e.processEnum(nest, sourceCodeInfo, nest_path)
	}

	for i, nest := range iter(msg.Messages()) {
		nest_path := append(path, 3, int32(i))
		if err := e.processMessage(nest, sourceCodeInfo, nest_path); err != nil {
			return err
		}
	}

	def := Message{
		Name:   resolveName(msg),
		Fields: []Field{},
	}
	def.LeadingComments, def.TrailingComments = extractComments(sourceCodeInfo, path)

	for i, field := range iter(msg.Fields()) {
		// fp := msgp.GetField()[i]
		typ, err := e.resolveType(def.Name, field)
		if err != nil {
			return fmt.Errorf("failed to resolve type: %w", err)
		}
		fieldpath := append(path, 2, int32(i))
		var oneOf *OneOf
		if oo := field.ContainingOneof(); !field.HasOptionalKeyword() && oo != nil {
			var fieldNames []string
			for _, f := range iter(oo.Fields()) {
				fieldNames = append(fieldNames, string(f.Name()))
			}
			oneOf = &OneOf{
				Name:       string(oo.Name()),
				FieldNames: fieldNames,
			}
		}
		name := field.JSONName()
		if e.config.PreservingProtoFieldName {
			name = string(field.Name())
		}
		f := Field{
			Name:     name,
			Type:     typ,
			Optional: field.HasOptionalKeyword(),
			OneOf:    oneOf,
			// Description: resolveFieldDescription(field),
		}
		f.LeadingComments, f.TrailingComments = extractComments(sourceCodeInfo, fieldpath)
		def.Fields = append(def.Fields, f)
	}

	e.messages = append(e.messages, def)
	return nil
}

func (e *generator) addExternalImport(importLine string) {
	for _, imp := range e.externalImports {
		if imp == importLine {
			return
		}
	}
	e.externalImports = append(e.externalImports, importLine)
}

// func (e *generator) depend(referer string, dependee string) {
// 	_, ok := e.deps[referer]
// 	if !ok {
// 		e.deps[referer] = []string{}
// 	}
// 	for _, dep := range e.deps[referer] {
// 		if dep == dependee {
// 			return
// 		}
// 	}
// 	e.deps[referer] = append(e.deps[referer], dependee)
// }

func (e *generator) resolveBaseType(referer string, field protoreflect.FieldDescriptor) (string, error) {
	switch field.Kind() {
	case
		protoreflect.Int32Kind,
		protoreflect.Int64Kind,
		protoreflect.Uint32Kind,
		protoreflect.Uint64Kind,
		protoreflect.Fixed32Kind,
		protoreflect.Fixed64Kind,
		protoreflect.Sint32Kind,
		protoreflect.Sint64Kind,
		protoreflect.Sfixed32Kind,
		protoreflect.Sfixed64Kind:
		return "int", nil
	case protoreflect.BoolKind:
		return "bool", nil
	case protoreflect.DoubleKind,
		protoreflect.FloatKind:
		return "float", nil
	case protoreflect.StringKind:
		return "str", nil
	case protoreflect.BytesKind:
		return "bytes", nil
	case protoreflect.MessageKind:
	case protoreflect.EnumKind:
		return resolveName(field.Enum()), nil
	case protoreflect.GroupKind:
		return "", fmt.Errorf("unsupported type: %s", field.Kind())
	}

	// Handle message types.
	msg := field.Message()
	name := resolveName(msg)

	if field.IsMap() {
		key, err := e.resolveType(referer, field.MapKey())
		if err != nil {
			return "", err
		}
		val, err := e.resolveType(referer, field.MapValue())
		if err != nil {
			return "", err
		}
		return fmt.Sprintf("dict[%s, %s]", key, val), nil
	}

	pkg := string(field.ParentFile().Package())
	fieldpkg := string(msg.ParentFile().Package())
	if pkg == fieldpkg {
		typ := resolveName(msg)
		// e.depend(referer, typ)
		return typ, nil
	}
	srcname := strings.TrimSuffix(filepath.Base(msg.ParentFile().Path()), ".proto")
	pkg = fmt.Sprintf("%s.%s_pb2", fieldpkg, srcname)
	fullname := fmt.Sprintf("%s.%s", pkg, name)
	e.addExternalImport(fmt.Sprintf("import %s", pkg))
	return fullname, nil
}

func (e *generator) resolveType(referer string, field protoreflect.FieldDescriptor) (string, error) {
	typ, err := e.resolveBaseType(referer, field)
	if err != nil {
		return "", err
	}

	if field.IsList() {
		return fmt.Sprintf("list[%s]", typ), nil
	}

	if field.HasOptionalKeyword() || field.ContainingOneof() != nil {
		if e.config.UseNoneUnionSyntaxInstaedOfOptional {
			return fmt.Sprintf("%s | None", typ), nil
		}
		return fmt.Sprintf("_Optional[%s]", typ), nil
	}

	return typ, nil
}

func equalPath(a, b []int32) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func resolveName(d protoreflect.Descriptor) string {
	prefix := string(d.ParentFile().FullName()) + "."
	name := string(d.FullName())
	name = strings.TrimPrefix(name, prefix)
	name = strings.ReplaceAll(name, ".", "_")
	return string(name)
}

func extractComments(sourceCodeInfo *descriptorpb.SourceCodeInfo, path []int32) (leading []string, trailing []string) {
	if sourceCodeInfo != nil {
		for _, location := range sourceCodeInfo.Location {
			if equalPath(location.Path, path) {
				leading = extractCommentLines(location.GetLeadingComments())
				trailing = extractCommentLines(location.GetTrailingComments())
				break
			}
		}
	}
	return
}

func extractCommentLines(comment string) []string {
	if comment == "" {
		return nil
	}
	comment = strings.TrimSpace(comment)
	comments := strings.Split(comment, "\n")
	for i := range comments {
		comments[i] = strings.TrimSpace(comments[i])
	}
	return comments
}

func iter[T any](d interface {
	Len() int
	Get(int) T
},
) func(func(k int, v T) bool) {
	return func(yield func(k int, v T) bool) {
		for i := range d.Len() {
			yield(i, d.Get(i))
		}
	}
}

func camelToSnakeCase(str string) string {
	snake := matchFirstCap.ReplaceAllString(str, "${1}_${2}")
	snake = matchAllCap.ReplaceAllString(snake, "${1}_${2}")
	return strings.ToUpper(snake)
}

// type TopologicalSortable interface {
// 	TopoKey() string
// }
//
// // topologicalSort performs a topological sort on the provided data given the dependencies.
// // `data` is a slice of items to sort, and `dependencies` maps an item key to a list of keys it depends on.
// func topologicalSort[T TopologicalSortable](data []T, dependencies map[string][]string) ([]T, error) {
// 	// Build the graph and calculate in-degrees (number of incoming edges)
// 	inDegree := make(map[string]int)
// 	graph := make(map[string][]string)
//
// 	for _, item := range data {
// 		key := item.TopoKey()
// 		inDegree[key] = 0
// 		graph[key] = []string{}
// 	}
//
// 	for node, deps := range dependencies {
// 		for _, dep := range deps {
// 			graph[dep] = append(graph[dep], node)
// 			inDegree[node]++
// 		}
// 	}
//
// 	// Find all nodes with no incoming edges
// 	var zeroInDegreeQueue []string
// 	for node, count := range inDegree {
// 		if count == 0 {
// 			zeroInDegreeQueue = append(zeroInDegreeQueue, node)
// 		}
// 	}
//
// 	// Sort the zero-in-degree nodes to ensure determinism
// 	sort.Strings(zeroInDegreeQueue)
//
// 	var result []T
// 	visited := 0
//
// 	for len(zeroInDegreeQueue) > 0 {
// 		// Remove the node from the queue
// 		current := zeroInDegreeQueue[0]
// 		zeroInDegreeQueue = zeroInDegreeQueue[1:]
//
// 		// Find the actual item corresponding to the current key
// 		var currentItem T
// 		for _, item := range data {
// 			if item.TopoKey() == current {
// 				currentItem = item
// 				break
// 			}
// 		}
//
// 		result = append(result, currentItem)
// 		visited++
//
// 		for _, neighbor := range graph[current] {
// 			inDegree[neighbor]--
// 			if inDegree[neighbor] == 0 {
// 				zeroInDegreeQueue = append(zeroInDegreeQueue, neighbor)
// 				sort.Strings(zeroInDegreeQueue) // Sort to maintain determinism
// 			}
// 		}
// 	}
//
// 	if visited == len(data) {
// 		return result, nil
// 	}
//
// 	return nil, errors.New("cycle detected in graph")
// }