github.com/stackb/rules_proto@v0.0.0-20240221195024-5428336c51f1/pkg/plugin/builtin/python_plugin.go (about)

     1  package builtin
     2  
     3  import (
     4  	"flag"
     5  	"log"
     6  	"path"
     7  	"path/filepath"
     8  	"strings"
     9  
    10  	"github.com/bazelbuild/bazel-gazelle/label"
    11  	"github.com/stackb/rules_proto/pkg/protoc"
    12  )
    13  
    14  func init() {
    15  	protoc.Plugins().MustRegisterPlugin(&PythonPlugin{})
    16  }
    17  
    18  // PythonPlugin implements Plugin for the built-in protoc python plugin.
    19  type PythonPlugin struct{}
    20  
    21  // Name implements part of the Plugin interface.
    22  func (p *PythonPlugin) Name() string {
    23  	return "builtin:python"
    24  }
    25  
    26  // Configure implements part of the Plugin interface.
    27  func (p *PythonPlugin) Configure(ctx *protoc.PluginContext) *protoc.PluginConfiguration {
    28  	flags := parsePythonPluginOptions(p.Name(), ctx.PluginConfig.GetFlags())
    29  
    30  	pyFiles := protoc.FlatMapFiles(
    31  		pythonGeneratedFileName(ctx.Rel),
    32  		protoc.Always,
    33  		ctx.ProtoLibrary.Files()...,
    34  	)
    35  
    36  	pyOutputs := make([]string, 0, len(pyFiles))
    37  	for _, pyFile := range pyFiles {
    38  		if flags.excludeOutput[filepath.Base(pyFile)] {
    39  			continue
    40  		}
    41  		pyOutputs = append(pyOutputs, pyFile)
    42  	}
    43  
    44  	return &protoc.PluginConfiguration{
    45  		Label:   label.New("build_stack_rules_proto", "plugin/builtin", "python"),
    46  		Outputs: pyOutputs,
    47  		Options: ctx.PluginConfig.GetOptions(),
    48  	}
    49  }
    50  
    51  // pythonGeneratedFileName is a utility function that returns a function that
    52  // computes the name of a predicted generated file having the given
    53  // extension(s) relative to the given dir.
    54  func pythonGeneratedFileName(reldir string) func(f *protoc.File) []string {
    55  	return func(f *protoc.File) []string {
    56  		name := strings.ReplaceAll(f.Name, "-", "_")
    57  		if reldir != "" {
    58  			name = path.Join(reldir, name)
    59  		}
    60  		return []string{name + "_pb2.py"}
    61  	}
    62  }
    63  
    64  // pythonPluginOptions represents the parsed flag configuration for the
    65  // ProtocGenTsProto implementation.
    66  type pythonPluginOptions struct {
    67  	excludeOutput map[string]bool
    68  }
    69  
    70  func parsePythonPluginOptions(kindName string, args []string) *pythonPluginOptions {
    71  	flags := flag.NewFlagSet(kindName, flag.ExitOnError)
    72  
    73  	var excludeOutput string
    74  	flags.StringVar(&excludeOutput, "exclude_output", "", "--exclude_output=foo_pb2.py suppresses the file 'foo_pb2.py' from the output list")
    75  
    76  	if err := flags.Parse(args); err != nil {
    77  		log.Fatalf("failed to parse flags for %q: %v", kindName, err)
    78  	}
    79  	config := &pythonPluginOptions{
    80  		excludeOutput: make(map[string]bool),
    81  	}
    82  	for _, value := range strings.Split(excludeOutput, ",") {
    83  		config.excludeOutput[value] = true
    84  	}
    85  
    86  	return config
    87  }