From e27067b81b77344a92ef02ffb54e4ec7d883825e Mon Sep 17 00:00:00 2001
From: Vojtech Vitek <vojtech.vitek@golang.cz>
Date: Tue, 18 Jul 2023 20:24:46 +0200
Subject: [PATCH] Fix stdlib/custom imports

---
 imports.go.tmpl | 69 +++++++++++++++++++++++++++++--------------------
 struct.go.tmpl  |  3 ++-
 types.go.tmpl   |  6 ++---
 3 files changed, 46 insertions(+), 32 deletions(-)

diff --git a/imports.go.tmpl b/imports.go.tmpl
index 89d4a77..5d78e66 100644
--- a/imports.go.tmpl
+++ b/imports.go.tmpl
@@ -3,44 +3,57 @@
 {{- $opts := .Opts -}}
 
 {{- /* Map of import paths. */ -}}
-{{- $imports := dict -}}
+{{- $stdlibImports := dict -}}
+{{- set $stdlibImports "context" "" -}}
+{{- set $stdlibImports "errors" "" -}}
+{{- set $stdlibImports "encoding/json" "" -}}
+{{- set $stdlibImports "fmt" "" -}}
+{{- set $stdlibImports "io/ioutil" "" -}}
+{{- set $stdlibImports "net/http" "" -}}
+
 {{- if $opts.client}}
-	{{- set $imports "bytes" "" -}}
-	{{- set $imports "io" "" -}}
-	{{- set $imports "net/url" "" -}}
-{{- else -}}
-	{{- set $imports "strings" "" -}}
+	{{- set $stdlibImports "bytes" "" -}}
+	{{- set $stdlibImports "io" "" -}}
+	{{- set $stdlibImports "net/url" "" -}}
+{{- end -}}
+{{- if $opts.server}}
+	{{- set $stdlibImports "strings" "" -}}
 {{- end -}}
-{{- set $imports "context" "" -}}
-{{- set $imports "errors" "" -}}
-{{- set $imports "encoding/json" "" -}}
-{{- set $imports "fmt" "" -}}
-{{- set $imports "io/ioutil" "" -}}
-{{- set $imports "net/http" "" -}}
 
-{{- /*TODO*/ -}}
-{{- if false -}}
-{{- set $imports "time" "" -}}
+{{- /* Import "time" if there's at least one timestamp. */ -}}
+{{ if eq $opts.importTypesFrom "" }}
+	{{- range $_, $type := $types -}}
+		{{- range $_, $field := $type.Fields -}}
+			{{- if $field.Type -}}
+				{{- if eq $field.Type.Expr "timestamp" -}}
+					{{- set $stdlibImports "time" "" -}}
+				{{- end -}}
+			{{- end -}}
+		{{- end -}}
+	{{- end -}}
 {{- end -}}
 
-{{- range $import, $rename := $imports }}
+{{- /* Print stdlib imports. */ -}}
+{{- range $import, $rename := $stdlibImports }}
 	{{if ne $rename ""}}{{$rename}} {{end}}"{{$import}}"
 {{- end -}}
 
-{{- $imports = dict -}}
-{{ range $_, $type := $types -}}
-	{{- range $_, $field := $type.Fields -}}
-		{{- range $meta := $field.Meta -}}
-			{{- if exists $meta "go.type.import" -}}
-				{{- set $imports (get $meta "go.type.import") "" -}}
+{{ if ne $opts.importTypesFrom "" }}
+
+	"{{ $opts.importTypesFrom }}"
+{{- else }}
+	{{- /* Print custom type imports. */ -}}
+	{{ $imports := dict }}
+	{{ range $_, $type := $types -}}
+		{{- range $_, $field := $type.Fields -}}
+			{{- range $meta := $field.Meta -}}
+				{{- if exists $meta "go.type.import" -}}
+					{{- set $imports (get $meta "go.type.import") "" -}}
+				{{- end -}}
 			{{- end -}}
 		{{- end -}}
-	{{- end -}}
-{{- end }}
-{{ if ne $opts.importTypesFrom "" }}
-	"{{ $opts.importTypesFrom -}}"
-{{ else }}
-	{{ range $import, $rename := $imports }}
+	{{- end }}
+	{{- range $import, $rename := $imports }}
 	{{if ne $rename ""}}{{$rename}} {{end}}"{{$import}}"
 	{{- end -}}
 {{- end }}
diff --git a/struct.go.tmpl b/struct.go.tmpl
index f2873c1..b0c6850 100644
--- a/struct.go.tmpl
+++ b/struct.go.tmpl
@@ -3,6 +3,7 @@
 {{- $name := .Name -}}
 {{- $fields := .Fields -}}
 {{- $typeMap := .TypeMap -}}
+{{- $typePrefix := .TypePrefix -}}
 
 type {{$name}} struct {
 {{- range $_, $field := $fields -}}
@@ -29,7 +30,7 @@ type {{$name}} struct {
 			{{- end -}}
 		{{- end -}}
 	{{- end }}
-	{{$fieldName}} {{template "type" dict "Type" $field.Type "CustomType" $customType "Optional" $field.Optional "TypeMap" $typeMap}} `{{$jsonTag}}{{if len $structTags}} {{end}}{{join (sort $structTags) " "}}`
+	{{$fieldName}} {{template "type" dict "Type" $field.Type "CustomType" $customType "Optional" $field.Optional "TypeMap" $typeMap "TypePrefix" $typePrefix}} `{{$jsonTag}}{{if len $structTags}} {{end}}{{join (sort $structTags) " "}}`
 {{- end}}
 }
 {{- end }}
diff --git a/types.go.tmpl b/types.go.tmpl
index 0088933..05df40c 100644
--- a/types.go.tmpl
+++ b/types.go.tmpl
@@ -11,11 +11,11 @@
 {{ range $_, $type := $types -}}
 
 {{- if eq $type.Kind "enum" }}
-{{template "enum" dict "Name" $type.Name "Type" $type.Type "Fields" $type.Fields }}
+{{template "enum" dict "Name" $type.Name "Type" $type.Type "TypePrefix" $typePrefix "Fields" $type.Fields}}
 {{ end -}}
 
 {{- if eq $type.Kind "struct" }}
-{{template "struct" dict "Name" $type.Name "Fields" $type.Fields "TypeMap" $typeMap}}
+{{template "struct" dict "Name" $type.Name "TypeMap" $typeMap "TypePrefix" $typePrefix "Fields" $type.Fields}}
 {{ end -}}
 
 {{- end -}}
@@ -25,7 +25,7 @@
 {{ range $_, $service := $services}}
 type {{$service.Name}} interface {
 	{{- range $_, $method := $service.Methods}}
-	{{.Name}}(ctx context.Context{{range $_, $input := $method.Inputs}}, {{$input.Name}} {{template "type" dict "Type" $input.Type "Optional" $input.Optional "TypeMap" $typeMap "TypePrefix" $typePrefix}}{{end}}) {{if len .Outputs}}({{end}}{{range $i, $output := .Outputs}}{{template "type" dict "Type" $output.Type "Optional" $output.Optional "TypeMap" $typeMap "TypePrefix" $typePrefix}}{{if lt $i (len $method.Outputs)}}, {{end}}{{end}}error{{if len $method.Outputs}}){{end}}
+	{{.Name}}(ctx context.Context{{range $_, $input := $method.Inputs}}, {{$input.Name}} {{template "type" dict "Type" $input.Type "TypeMap" $typeMap "TypePrefix" $typePrefix "Optional" $input.Optional}}{{end}}) {{if len .Outputs}}({{end}}{{range $i, $output := .Outputs}}{{template "type" dict "Type" $output.Type "TypeMap" $typeMap "TypePrefix" $typePrefix "Optional" $output.Optional}}{{if lt $i (len $method.Outputs)}}, {{end}}{{end}}error{{if len $method.Outputs}}){{end}}
 	{{- end}}
 }
 {{end}}