Add extra imports

interfaces/generate.go

@@ -22,6 +22,9 @@ type Options struct {
 	// the method will be included. MethodHasPrefix and MethodHasSuffix can be used inside a custom function here
 	// to customize the behavior.
 	ShouldInclude func(reflect.Method) bool
+
+	// ExtraImports can add extra imports for a method
+	ExtraImports func(reflect.Method) []string
 }
 
 func (o *Options) SetDefaults() {
@@ -38,6 +41,12 @@ func WithIncludeFunc(f func(reflect.Method) bool) Option {
 	}
 }
 
+func WithExtraImports(f func(reflect.Method) []string) Option {
+	return func(o *Options) {
+		o.ExtraImports = f
+	}
+}
+
 // Generate generates service interfaces to be used for generating
 // mocks. The clients passed in as the first argument should be structs that will be used to
 // generate the service interfaces. The second argument, dir, is the path to the output
@@ -116,11 +125,12 @@ func signature(name string, f any) string {
 }
 
 type serviceInfo struct {
-	Import      string
-	Name        string
-	PackageName string
-	ClientName  string
-	Signatures  []string
+	Import       string
+	Name         string
+	PackageName  string
+	ClientName   string
+	Signatures   []string
+	ExtraImports []string
 }
 
 func getServiceInfo(client any, opts *Options) serviceInfo {
@@ -133,19 +143,22 @@ func getServiceInfo(client any, opts *Options) serviceInfo {
 	name := csr.ToPascal(pkgName)
 	clientName := name + "Client"
 	signatures := make([]string, 0)
+	extraImports := make([]string, 0)
 	for i := 0; i < t.NumMethod(); i++ {
 		method := t.Method(i)
 		if opts.ShouldInclude(method) {
 			sig := signature(method.Name, v.Method(i).Interface())
 			signatures = append(signatures, sig)
 		}
+		extraImports = append(extraImports, opts.ExtraImports(method)...)
 	}
 	return serviceInfo{
-		Import:      pkgPath,
-		Name:        name,
-		PackageName: pkgName,
-		ClientName:  clientName,
-		Signatures:  signatures,
+		Import:       pkgPath,
+		Name:         name,
+		PackageName:  pkgName,
+		ClientName:   clientName,
+		Signatures:   signatures,
+		ExtraImports: extraImports,
 	}
 }
 

interfaces/templates/service.go.tpl

@@ -2,8 +2,10 @@
 package services
 
 import (
-    "net/http"
     "{{ .Import }}"
+    {{- range .ExtraImports }}
+    "{{ . }}"
+    {{- end }}
 )
 
 //go:generate mockgen -package=mocks -destination=../mocks/{{.PackageName}}.go -source={{.PackageName}}.go {{.ClientName}}