Add buf generate --include-wkt (bufbuild#892)

CHANGELOG.md

@@ -5,6 +5,10 @@
 - Check that the user provided a valid token when running `buf registry login`.
 - Add `buf mod open` that opens a module's homepage in a browser.
 - Add `buf completion` command to generate auto-completion scripts in commonly used shells.
+- Add `--include-wkt` flag to `buf generate`. When this flag is specified alongside
+  `--include-imports`, this will result in the [Well-Known Types](https://github.com/bufbuild/wellknowntypes/tree/11ea259bf71c4d386131c1986ffe103cb1edb3d6/v3.19.4/google/protobuf)
+  being generated as well. Most language runtimes have the Well-Known Types included as part
+  of the core library, making generating the Well-Known Types separately undesirable.
 
 ## [v1.0.0-rc12] - 2022-02-01
 

private/buf/bufgen/bufgen.go

@@ -139,13 +139,23 @@ func GenerateWithBaseOutDirPath(baseOutDirPath string) GenerateOption {
 
 // GenerateWithIncludeImports says to also generate imports.
 //
-// Note that this does NOT result in the Well-Known Types being generated.
+// Note that this does NOT result in the Well-Known Types being generated, use
+// GenerateWithIncludeWellKnownTypes to include the Well-Known Types.
 func GenerateWithIncludeImports() GenerateOption {
 	return func(generateOptions *generateOptions) {
 		generateOptions.includeImports = true
 	}
 }
 
+// GenerateWithIncludeWellKnownTypes says to also generate well known types.
+//
+// This option has no effect if GenerateWithIncludeImports is not set.
+func GenerateWithIncludeWellKnownTypes() GenerateOption {
+	return func(generateOptions *generateOptions) {
+		generateOptions.includeWellKnownTypes = true
+	}
+}
+
 // Config is a configuration.
 type Config struct {
 	// Required

private/buf/bufgen/generator.go

@@ -92,6 +92,7 @@ func (g *generator) Generate(
 		image,
 		generateOptions.baseOutDirPath,
 		generateOptions.includeImports,
+		generateOptions.includeWellKnownTypes,
 	)
 }
 
@@ -102,6 +103,7 @@ func (g *generator) generate(
 	image bufimage.Image,
 	baseOutDirPath string,
 	includeImports bool,
+	includeWellKnownTypes bool,
 ) error {
 	if err := modifyImage(ctx, g.logger, config, image); err != nil {
 		return err
@@ -113,6 +115,7 @@ func (g *generator) generate(
 		config,
 		image,
 		includeImports,
+		includeWellKnownTypes,
 	)
 	if err != nil {
 		return err
@@ -153,6 +156,7 @@ func (g *generator) execPlugins(
 	config *Config,
 	image bufimage.Image,
 	includeImports bool,
+	includeWellKnownTypes bool,
 ) ([]*pluginpb.CodeGeneratorResponse, error) {
 	imageProvider := newImageProvider(image)
 	// Collect all of the plugin jobs so that they can be executed in parallel.
@@ -169,6 +173,7 @@ func (g *generator) execPlugins(
 					image,
 					currentPluginConfig,
 					includeImports,
+					includeWellKnownTypes,
 				)
 				if err != nil {
 					return err
@@ -186,6 +191,7 @@ func (g *generator) execPlugins(
 					imageProvider,
 					currentPluginConfig,
 					includeImports,
+					includeWellKnownTypes,
 				)
 				if err != nil {
 					return err
@@ -232,6 +238,7 @@ func (g *generator) execLocalPlugin(
 	imageProvider *imageProvider,
 	pluginConfig *PluginConfig,
 	includeImports bool,
+	includeWellKnownTypes bool,
 ) (*pluginpb.CodeGeneratorResponse, error) {
 	pluginImages, err := imageProvider.GetImages(pluginConfig.Strategy)
 	if err != nil {
@@ -246,6 +253,7 @@ func (g *generator) execLocalPlugin(
 			pluginConfig.Opt,
 			nil,
 			includeImports,
+			includeWellKnownTypes,
 		),
 		appprotoexec.GenerateWithPluginPath(pluginConfig.Path),
 	)
@@ -261,6 +269,7 @@ func (g *generator) execRemotePlugin(
 	image bufimage.Image,
 	pluginConfig *PluginConfig,
 	includeImports bool,
+	includeWellKnownTypes bool,
 ) (*pluginpb.CodeGeneratorResponse, error) {
 	remote, owner, name, version, err := bufplugin.ParsePluginVersionPath(pluginConfig.Remote)
 	if err != nil {
@@ -287,6 +296,7 @@ func (g *generator) execRemotePlugin(
 			},
 		},
 		includeImports,
+		includeWellKnownTypes,
 	)
 	if err != nil {
 		return nil, err
@@ -470,8 +480,9 @@ func validateResponses(
 }
 
 type generateOptions struct {
-	baseOutDirPath string
-	includeImports bool
+	baseOutDirPath        string
+	includeImports        bool
+	includeWellKnownTypes bool
 }
 
 func newGenerateOptions() *generateOptions {

private/buf/cmd/buf/command/generate/generate.go

@@ -40,6 +40,7 @@ const (
 	configFlagName              = "config"
 	pathsFlagName               = "path"
 	includeImportsFlagName      = "include-imports"
+	includeWKTFlagName          = "include-wkt"
 	excludePathsFlagName        = "exclude-path"
 
 	// deprecated
@@ -194,6 +195,7 @@ type flags struct {
 	Config         string
 	Paths          []string
 	IncludeImports bool
+	IncludeWKT     bool
 	ExcludePaths   []string
 
 	// deprecated
@@ -218,6 +220,15 @@ func (f *flags) Bind(flagSet *pflag.FlagSet) {
 		false,
 		"Also generate all imports except for Well-Known Types.",
 	)
+	flagSet.BoolVar(
+		&f.IncludeWKT,
+		includeWKTFlagName,
+		false,
+		fmt.Sprintf(
+			"Also generate Well-Known Types. Cannot be set without --%s.",
+			includeImportsFlagName,
+		),
+	)
 	flagSet.StringVar(
 		&f.Template,
 		templateFlagName,
@@ -274,6 +285,13 @@ func run(
 	flags *flags,
 ) (retErr error) {
 	logger := container.Logger()
+	if flags.IncludeWKT && !flags.IncludeImports {
+		// You need to set --include-imports if you set --include-wkt, which isn’t great. The alternative is to have
+		// --include-wkt implicitly set --include-imports, but this could be surprising. Or we could rename
+		// --include-wkt to --include-imports-and/with-wkt. But the summary is that the flag only makes sense
+		// in the context of including imports.
+		return appcmd.NewInvalidArgumentErrorf("Cannot set --%s without --%s", includeWKTFlagName, includeImportsFlagName)
+	}
 	if err := bufcli.ValidateErrorFormatFlag(flags.ErrorFormat, errorFormatFlagName); err != nil {
 		return err
 	}
@@ -370,6 +388,12 @@ func run(
 			bufgen.GenerateWithIncludeImports(),
 		)
 	}
+	if flags.IncludeWKT {
+		generateOptions = append(
+			generateOptions,
+			bufgen.GenerateWithIncludeWellKnownTypes(),
+		)
+	}
 	return bufgen.NewGenerator(
 		logger,
 		storageosProvider,

private/buf/cmd/buf/command/protoc/plugin.go

@@ -64,6 +64,7 @@ func executePlugin(
 			strings.Join(pluginInfo.Opt, ","),
 			appprotoexec.DefaultVersion,
 			false,
+			false,
 		),
 		appprotoexec.GenerateWithPluginPath(pluginInfo.Path),
 	)

private/buf/cmd/protoc-gen-buf-lint/lint_test.go

@@ -289,5 +289,5 @@ func testBuildCodeGeneratorRequest(
 	}
 	image, err := bufimage.NewImage(imageFiles)
 	require.NoError(t, err)
-	return bufimage.ImageToCodeGeneratorRequest(image, parameter, nil, false)
+	return bufimage.ImageToCodeGeneratorRequest(image, parameter, nil, false, false)
 }

private/bufpkg/bufimage/bufimage.go

@@ -362,17 +362,21 @@ func ImageToFileDescriptorProtos(image Image) []*descriptorpb.FileDescriptorProt
 //
 // All non-imports are added as files to generate.
 // If includeImports is set, all non-well-known-type imports are also added as files to generate.
+// If includeWellKnownTypes is set, well-known-type imports are also added as files to generate.
+// includeWellKnownTypes has no effect if includeImports is not set.
 func ImageToCodeGeneratorRequest(
 	image Image,
 	parameter string,
 	compilerVersion *pluginpb.Version,
 	includeImports bool,
+	includeWellKnownTypes bool,
 ) *pluginpb.CodeGeneratorRequest {
 	return imageToCodeGeneratorRequest(
 		image,
 		parameter,
 		compilerVersion,
 		includeImports,
+		includeWellKnownTypes,
 		nil,
 		nil,
 	)
@@ -383,17 +387,37 @@ func ImageToCodeGeneratorRequest(
 // All non-imports are added as files to generate.
 // If includeImports is set, all non-well-known-type imports are also added as files to generate.
 // If includeImports is set, only one CodeGeneratorRequest will contain any given file as a FileToGenerate.
+// If includeWellKnownTypes is set, well-known-type imports are also added as files to generate.
+// includeWellKnownTypes has no effect if includeImports is not set.
 func ImagesToCodeGeneratorRequests(
 	images []Image,
 	parameter string,
 	compilerVersion *pluginpb.Version,
 	includeImports bool,
+	includeWellKnownTypes bool,
 ) []*pluginpb.CodeGeneratorRequest {
 	requests := make([]*pluginpb.CodeGeneratorRequest, len(images))
-	// we don't need to track these if includeImports as false, so don't waste the time
+	// alreadyUsedPaths is a map of paths that have already been added to an image.
+	//
+	// We track this if includeImports is set, so that when we find an import, we can
+	// see if the import was already added to a CodeGeneratorRequest via another Image
+	// in the Image slice. If the import was already added, we do not add duplicates
+	// across CodeGeneratorRequests.
 	var alreadyUsedPaths map[string]struct{}
+	// nonImportPaths is a map of non-import paths.
+	//
+	// We track this if includeImports is set. If we find a non-import file in Image A
+	// and this file is an import in Image B, the file will have already been added to
+	// a CodeGeneratorRequest via Image A, so do not add the duplicate to any other
+	// CodeGeneratorRequest.
 	var nonImportPaths map[string]struct{}
 	if includeImports {
+		// We don't need to track these if includeImports is false, so we only populate
+		// the maps if includeImports is true. If includeImports is false, only non-imports
+		// will be added to each CodeGeneratorRequest, so figuring out whether or not
+		// we should add a given import to a given CodeGeneratorRequest is unnecessary.
+		//
+		// imageToCodeGeneratorRequest checks if these maps are nil before every access.
 		alreadyUsedPaths = make(map[string]struct{})
 		nonImportPaths = make(map[string]struct{})
 		for _, image := range images {
@@ -410,6 +434,7 @@ func ImagesToCodeGeneratorRequests(
 			parameter,
 			compilerVersion,
 			includeImports,
+			includeWellKnownTypes,
 			alreadyUsedPaths,
 			nonImportPaths,
 		)

private/bufpkg/bufimage/bufimagemodify/bufimagemodify.go

@@ -282,10 +282,7 @@ func RubyPackage(
 
 // isWellKnownType returns true if the given path is one of the well-known types.
 func isWellKnownType(ctx context.Context, imageFile bufimage.ImageFile) bool {
-	if _, err := datawkt.ReadBucket.Stat(ctx, imageFile.Path()); err == nil {
-		return true
-	}
-	return false
+	return datawkt.Exists(imageFile.Path())
 }
 
 // int32SliceIsEqual returns true if x and y contain the same elements.