feat: supports init mode

Bin-Huang · Aug 4, 2022 · 305363a · 305363a
1 parent c2b9e17
commit 305363a
Show file tree

Hide file tree

Showing 7 changed files with 44 additions and 10 deletions.
diff --git a/generate.go b/generate.go
@@ -15,21 +15,29 @@ var templ = `// Code generated by github.com/Bin-Huang/make-constructor; DO NOT
 package {{.PkgName}}
 
 import (
-{{ range .Imports }}
-    {{ if .Name }}
+{{ range .Imports -}}
+    {{ if .Name -}}
     {{.Name}} {{.Path}}
-    {{ else }}
+    {{ else -}}
     {{.Path}}
-    {{ end }}
-{{ end }}
+    {{ end -}}
+{{ end -}}
 )
 
 {{ range .Constructors }}
 // {{.Name}} Create a new {{.Struct}}
 func {{.Name}}({{.Params}}) *{{.Struct}} {
+	{{ if .Init -}}
+    s := &{{.Struct}} {
+        {{.Fields}}
+    }
+	s.init()
+	return s
+	{{ else -}}
     return &{{.Struct}} {
         {{.Fields}}
     }
+	{{ end -}}
 }
 {{ end }}
 `
@@ -59,6 +67,7 @@ func generateCode(pkgName string, importResuts []ResultImport, results []Result)
 		constructors = append(constructors, o{
 			"Name":   "New" + strcase.ToCamel(result.StructName),
 			"Struct": result.StructName,
+			"Init":   result.Init,
 			"Params": strings.Join(params, ", "),
 			"Fields": strings.Join(fields, "\n"),
 		})

diff --git a/main.go b/main.go
@@ -98,6 +98,11 @@ func IsMakeComment(s string) bool {
 	return strings.HasPrefix(s, "//go:generate") && strings.Contains(s, "make-constructor")
 }
 
+// IsInitModeEnable check if this struct enable the init mode
+func IsInitModeEnable(s string) bool {
+	return strings.Contains(s, "init")
+}
+
 // BuildAST ...
 func BuildAST(filename string) (*ast.File, error) {
 	astFile, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
@@ -122,6 +127,7 @@ type ResultImport struct {
 // Result ...
 type Result struct {
 	StructName string
+	Init       bool
 	Fields     []ResultField
 }
 
@@ -139,11 +145,13 @@ func ParseFile(filename string) ([]Result, []ResultImport, error) {
 			continue
 		}
 
+		var initMode bool
 		if genDecl.Tok == token.TYPE {
 			needGen := false
 			for _, doc := range genDecl.Doc.List {
 				if IsMakeComment(doc.Text) {
 					needGen = true
+					initMode = IsInitModeEnable(doc.Text)
 					break
 				}
 			}
@@ -201,6 +209,7 @@ func ParseFile(filename string) ([]Result, []ResultImport, error) {
 			results = append(results, Result{
 				StructName: typeSpec.Name.Name,
 				Fields:     resultFields,
+				Init:       initMode,
 			})
 		}
 	}

diff --git a/test.sh b/test.sh
@@ -2,3 +2,4 @@ set -ex
 
 go generate ./...
 go vet ./...
+go test ./...
diff --git a/test/repositories/constructor_gen.go b/test/repositories/constructor_gen.go
diff --git a/test/repositories/init_mode_test.go b/test/repositories/init_mode_test.go
@@ -0,0 +1,10 @@
+package repositories
+
+import "testing"
+
+func TestInitMode(t *testing.T) {
+	r := NewUserRepository(nil, nil, "")
+	if r.TableName != "foo" {
+		t.Error("generated init-mode code invalid: the constructor of UserRepository should calls `s.init()` method")
+	}
+}
diff --git a/test/repositories/user-repository.go b/test/repositories/user-repository.go
@@ -7,13 +7,17 @@ import (
 )
 
 // UserRepository the user repository for example
-//go:generate go run ../../../make-constructor
+//go:generate go run ../../../make-constructor --init
 type UserRepository struct {
 	conf      *config2.Config
 	db        *database
 	TableName string
 }
 
+func (r *UserRepository) init() {
+	r.TableName = "foo"
+}
+
 // FindByID find something by id
 func (r *UserRepository) FindByID() error {
 	return errors.New("no found")

diff --git a/test/services/constructor_gen.go b/test/services/constructor_gen.go