From 8be7c15c80f5358bb0332b9b9e33944870806357 Mon Sep 17 00:00:00 2001
From: Pantani <Pantani>
Date: Wed, 22 May 2024 01:10:37 +0200
Subject: [PATCH] add include import optional

---
 ignite/pkg/cosmosbuf/buf.go                 | 36 ++++++++++++++++-----
 ignite/pkg/cosmosgen/generate_go.go         |  9 +++++-
 ignite/pkg/cosmosgen/generate_typescript.go |  1 +
 3 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/ignite/pkg/cosmosbuf/buf.go b/ignite/pkg/cosmosbuf/buf.go
index 91e2221a73..d7ca39af9e 100644
--- a/ignite/pkg/cosmosbuf/buf.go
+++ b/ignite/pkg/cosmosbuf/buf.go
@@ -61,9 +61,11 @@ type (
 
 	// genOptions used to configure code generation.
 	genOptions struct {
-		excluded   []glob.Glob
-		fileByFile bool
-		flags      map[string]string
+		excluded       []glob.Glob
+		fileByFile     bool
+		flags          map[string]string
+		includeImports string
+		includeWKT     string
 	}
 
 	// GenOption configures code generation.
@@ -72,9 +74,11 @@ type (
 
 func newGenOptions() genOptions {
 	return genOptions{
-		flags:      make(map[string]string),
-		excluded:   make([]glob.Glob, 0),
-		fileByFile: false,
+		flags:          make(map[string]string),
+		excluded:       make([]glob.Glob, 0),
+		fileByFile:     false,
+		includeWKT:     "false",
+		includeImports: "false",
 	}
 }
 
@@ -94,6 +98,22 @@ func ExcludeFiles(patterns ...string) GenOption {
 	}
 }
 
+// IncludeImports also generate all imports except for Well-Known Types
+func IncludeImports() GenOption {
+	return func(o *genOptions) {
+		o.includeImports = "true"
+	}
+}
+
+// IncludeWKT also generate Well-Known Types.
+// Cannot be set to true without setting --include-imports to true
+func IncludeWKT() GenOption {
+	return func(o *genOptions) {
+		o.includeImports = "true"
+		o.includeWKT = "true"
+	}
+}
+
 // FileByFile runs the generate command for each proto file.
 func FileByFile() GenOption {
 	return func(o *genOptions) {
@@ -206,8 +226,8 @@ func (b Buf) Generate(
 		flagOutput:                output,
 		flagErrorFormat:           fmtJSON,
 		flagLogFormat:             fmtJSON,
-		flagIncludeImports:        "true",
-		flagIncludeWellKnownTypes: "true",
+		flagIncludeImports:        opts.includeImports,
+		flagIncludeWellKnownTypes: opts.includeWKT,
 	}
 	for k, v := range opts.flags {
 		flags[k] = v
diff --git a/ignite/pkg/cosmosgen/generate_go.go b/ignite/pkg/cosmosgen/generate_go.go
index 2aeb52f6c9..781306c2f5 100644
--- a/ignite/pkg/cosmosgen/generate_go.go
+++ b/ignite/pkg/cosmosgen/generate_go.go
@@ -42,7 +42,14 @@ func (g *generator) generate(ctx context.Context, template, fromPath string, exc
 	defer os.RemoveAll(tmp)
 
 	// code generate for each module.
-	if err := g.buf.Generate(ctx, g.protoPath(), tmp, template, cosmosbuf.ExcludeFiles(excluded...)); err != nil {
+	if err := g.buf.Generate(
+		ctx,
+		g.protoPath(),
+		tmp,
+		template,
+		cosmosbuf.ExcludeFiles(excluded...),
+		cosmosbuf.IncludeWKT(),
+	); err != nil {
 		return err
 	}
 
diff --git a/ignite/pkg/cosmosgen/generate_typescript.go b/ignite/pkg/cosmosgen/generate_typescript.go
index bed25f2f0e..8ed6893350 100644
--- a/ignite/pkg/cosmosgen/generate_typescript.go
+++ b/ignite/pkg/cosmosgen/generate_typescript.go
@@ -154,6 +154,7 @@ func (g *tsGenerator) generateModuleTemplate(
 		typesOut,
 		g.g.tsTemplate(),
 		cosmosbuf.ExcludeFiles("module.proto"),
+		cosmosbuf.IncludeWKT(),
 	); err != nil {
 		return err
 	}