Merge pull request #762 from flycash/generateRouters

Add bee generate routers command

cmd/commands/generate/generate.go

@@ -53,6 +53,10 @@ var CmdGenerate = &commands.Command{
 
      $ bee generate docs
 
+    ▶ {{"To generate swagger doc file:"|bold}}
+
+     $ bee generate routers [-ctrlDir=/path/to/controller/directory] [-routersFile=/path/to/routers/file.go] [-routersPkg=myPackage]
+
   ▶ {{"To generate a test case:"|bold}}
 
      $ bee generate test [routerfile]
@@ -72,6 +76,15 @@ func init() {
 	CmdGenerate.Flag.Var(&generate.Level, "level", "Either 1, 2 or 3. i.e. 1=models; 2=models and controllers; 3=models, controllers and routers.")
 	CmdGenerate.Flag.Var(&generate.Fields, "fields", "List of table Fields.")
 	CmdGenerate.Flag.Var(&generate.DDL, "ddl", "Generate DDL Migration")
+
+	// bee generate routers
+	CmdGenerate.Flag.Var(&generate.ControllerDirectory, "ctrlDir",
+		"Controller directory. Bee scans this directory and its sub directory to generate routers")
+	CmdGenerate.Flag.Var(&generate.RoutersFile, "routersFile",
+		"Routers file. If not found, Bee create a new one. Bee will truncates this file and output routers info into this file")
+	CmdGenerate.Flag.Var(&generate.RouterPkg, "routersPkg",
+		`router's package. Default is routers, it means that "package routers" in the generated file`)
+
 	commands.AvailableCommands = append(commands.AvailableCommands, CmdGenerate)
 }
 
@@ -97,6 +110,8 @@ func GenerateCode(cmd *commands.Command, args []string) int {
 		model(cmd, args, currpath)
 	case "view":
 		view(args, currpath)
+	case "routers":
+		genRouters(cmd, args)
 	default:
 		beeLogger.Log.Fatal("Command is missing")
 	}
@@ -104,6 +119,16 @@ func GenerateCode(cmd *commands.Command, args []string) int {
 	return 0
 }
 
+func genRouters(cmd *commands.Command, args []string) {
+	err := cmd.Flag.Parse(args[1:])
+	beeLogger.Log.Infof("input parameter: %v", args)
+	if err != nil {
+		beeLogger.Log.Errorf("could not parse input parameter: %+v", err)
+		return
+	}
+	generate.GenRouters()
+}
+
 func scaffold(cmd *commands.Command, args []string, currpath string) {
 	if len(args) < 2 {
 		beeLogger.Log.Fatal("Wrong number of arguments. Run: bee help generate")

generate/g.go

@@ -22,3 +22,9 @@ var Level utils.DocValue
 var Tables utils.DocValue
 var Fields utils.DocValue
 var DDL utils.DocValue
+
+
+// bee generate routers
+var ControllerDirectory utils.DocValue
+var RoutersFile utils.DocValue
+var RouterPkg utils.DocValue

generate/gen_routes.go

@@ -0,0 +1,547 @@
+// Copyright 2014 beego Author. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package generate
+
+import (
+	"errors"
+	"fmt"
+	"go/ast"
+	"os"
+	"path/filepath"
+	"regexp"
+	"sort"
+	"strconv"
+	"strings"
+	"unicode"
+
+	"github.com/beego/beego/v2/server/web"
+	"golang.org/x/tools/go/packages"
+
+	"github.com/beego/beego/v2/core/logs"
+
+	"github.com/beego/beego/v2/server/web/context/param"
+
+	beeLogger "github.com/beego/bee/v2/logger"
+)
+
+var globalRouterTemplate = `package {{.routersDir}}
+
+import (
+	beego "github.com/beego/beego/v2/server/web"
+	"github.com/beego/beego/v2/server/web/context/param"{{.globalimport}}
+)
+
+func init() {
+{{.globalinfo}}
+}
+`
+
+var (
+	genInfoList map[string][]web.ControllerComments
+
+	routerHooks = map[string]int{
+		"beego.BeforeStatic": web.BeforeStatic,
+		"beego.BeforeRouter": web.BeforeRouter,
+		"beego.BeforeExec":   web.BeforeExec,
+		"beego.AfterExec":    web.AfterExec,
+		"beego.FinishRouter": web.FinishRouter,
+	}
+
+	routerHooksMapping = map[int]string{
+		web.BeforeStatic: "beego.BeforeStatic",
+		web.BeforeRouter: "beego.BeforeRouter",
+		web.BeforeExec:   "beego.BeforeExec",
+		web.AfterExec:    "beego.AfterExec",
+		web.FinishRouter: "beego.FinishRouter",
+	}
+)
+
+
+
+func GenRouters() {
+	ctrlPath := ControllerDirectory.String()
+	if len(ctrlPath) == 0 {
+		ctrlPath = "controllers"
+	}
+	routersPath := RoutersFile.String()
+	if len(routersPath) == 0 {
+		routersPath = "routers/commentsRouter.go"
+	}
+	genRouters(ctrlPath, routersPath)
+}
+
+func genRouters(ctrlPath string, routersPath string) {
+	beeLogger.Log.Infof("read controller info from directory[%s], and output routers to [%s]", ctrlPath, routersPath)
+	err :=  parserPkg(ctrlPath, routersPath)
+	if err != nil {
+		beeLogger.Log.Errorf("Could not generate routers. Please check your input: %+v", err)
+	}
+}
+
+func parserPkg(ctrlPath string, routersPath string) error {
+	genInfoList = make(map[string][]web.ControllerComments)
+	pkgs, err := packages.Load(&packages.Config{
+		Mode: packages.NeedName | packages.NeedFiles | packages.NeedCompiledGoFiles | packages.NeedSyntax,
+		Dir:  ctrlPath,
+	}, "./...")
+
+	if err != nil {
+		return err
+	}
+	for _, pkg := range pkgs {
+		for _, fl := range pkg.Syntax {
+			for _, d := range fl.Decls {
+				switch specDecl := d.(type) {
+				case *ast.FuncDecl:
+					if specDecl.Recv != nil {
+						exp, ok := specDecl.Recv.List[0].Type.(*ast.StarExpr) // Check that the type is correct first beforing throwing to parser
+						if ok {
+							err = parserComments(specDecl, fmt.Sprint(exp.X), pkg.PkgPath)
+							if err != nil {
+								return err
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+	return genRouterCode(routersPath)
+}
+
+type parsedComment struct {
+	routerPath string
+	methods    []string
+	params     map[string]parsedParam
+	filters    []parsedFilter
+	imports    []parsedImport
+}
+
+type parsedImport struct {
+	importPath  string
+	importAlias string
+}
+
+type parsedFilter struct {
+	pattern string
+	pos     int
+	filter  string
+	params  []bool
+}
+
+type parsedParam struct {
+	name     string
+	datatype string
+	location string
+	defValue string
+	required bool
+}
+
+func parserComments(f *ast.FuncDecl, controllerName, pkgpath string) error {
+	if f.Doc != nil {
+		parsedComments, err := parseComment(f.Doc.List)
+		if err != nil {
+			return err
+		}
+		for _, parsedComment := range parsedComments {
+			if parsedComment.routerPath != "" {
+				key := pkgpath + ":" + controllerName
+				cc := web.ControllerComments{}
+				cc.Method = f.Name.String()
+				cc.Router = parsedComment.routerPath
+				cc.AllowHTTPMethods = parsedComment.methods
+				cc.MethodParams = buildMethodParams(f.Type.Params.List, parsedComment)
+				cc.FilterComments = buildFilters(parsedComment.filters)
+				cc.ImportComments = buildImports(parsedComment.imports)
+				genInfoList[key] = append(genInfoList[key], cc)
+			}
+		}
+	}
+	return nil
+}
+
+func buildImports(pis []parsedImport) []*web.ControllerImportComments {
+	var importComments []*web.ControllerImportComments
+
+	for _, pi := range pis {
+		importComments = append(importComments, &web.ControllerImportComments{
+			ImportPath:  pi.importPath,
+			ImportAlias: pi.importAlias,
+		})
+	}
+
+	return importComments
+}
+
+func buildFilters(pfs []parsedFilter) []*web.ControllerFilterComments {
+	var filterComments []*web.ControllerFilterComments
+
+	for _, pf := range pfs {
+		var (
+			returnOnOutput bool
+			resetParams    bool
+		)
+
+		if len(pf.params) >= 1 {
+			returnOnOutput = pf.params[0]
+		}
+
+		if len(pf.params) >= 2 {
+			resetParams = pf.params[1]
+		}
+
+		filterComments = append(filterComments, &web.ControllerFilterComments{
+			Filter:         pf.filter,
+			Pattern:        pf.pattern,
+			Pos:            pf.pos,
+			ReturnOnOutput: returnOnOutput,
+			ResetParams:    resetParams,
+		})
+	}
+
+	return filterComments
+}
+
+func buildMethodParams(funcParams []*ast.Field, pc *parsedComment) []*param.MethodParam {
+	result := make([]*param.MethodParam, 0, len(funcParams))
+	for _, fparam := range funcParams {
+		for _, pName := range fparam.Names {
+			methodParam := buildMethodParam(fparam, pName.Name, pc)
+			result = append(result, methodParam)
+		}
+	}
+	return result
+}
+
+func buildMethodParam(fparam *ast.Field, name string, pc *parsedComment) *param.MethodParam {
+	options := []param.MethodParamOption{}
+	if cparam, ok := pc.params[name]; ok {
+		// Build param from comment info
+		name = cparam.name
+		if cparam.required {
+			options = append(options, param.IsRequired)
+		}
+		switch cparam.location {
+		case "body":
+			options = append(options, param.InBody)
+		case "header":
+			options = append(options, param.InHeader)
+		case "path":
+			options = append(options, param.InPath)
+		}
+		if cparam.defValue != "" {
+			options = append(options, param.Default(cparam.defValue))
+		}
+	} else {
+		if paramInPath(name, pc.routerPath) {
+			options = append(options, param.InPath)
+		}
+	}
+	return param.New(name, options...)
+}
+
+func paramInPath(name, route string) bool {
+	return strings.HasSuffix(route, ":"+name) ||
+		strings.Contains(route, ":"+name+"/")
+}
+
+var routeRegex = regexp.MustCompile(`@router\s+(\S+)(?:\s+\[(\S+)\])?`)
+
+func parseComment(lines []*ast.Comment) (pcs []*parsedComment, err error) {
+	pcs = []*parsedComment{}
+	params := map[string]parsedParam{}
+	filters := []parsedFilter{}
+	imports := []parsedImport{}
+
+	for _, c := range lines {
+		t := strings.TrimSpace(strings.TrimLeft(c.Text, "//"))
+		if strings.HasPrefix(t, "@Param") {
+			pv := getparams(strings.TrimSpace(strings.TrimLeft(t, "@Param")))
+			if len(pv) < 4 {
+				logs.Error("Invalid @Param format. Needs at least 4 parameters")
+			}
+			p := parsedParam{}
+			names := strings.SplitN(pv[0], "=>", 2)
+			p.name = names[0]
+			funcParamName := p.name
+			if len(names) > 1 {
+				funcParamName = names[1]
+			}
+			p.location = pv[1]
+			p.datatype = pv[2]
+			switch len(pv) {
+			case 5:
+				p.required, _ = strconv.ParseBool(pv[3])
+			case 6:
+				p.defValue = pv[3]
+				p.required, _ = strconv.ParseBool(pv[4])
+			}
+			params[funcParamName] = p
+		}
+	}
+
+	for _, c := range lines {
+		t := strings.TrimSpace(strings.TrimLeft(c.Text, "//"))
+		if strings.HasPrefix(t, "@Import") {
+			iv := getparams(strings.TrimSpace(strings.TrimLeft(t, "@Import")))
+			if len(iv) == 0 || len(iv) > 2 {
+				logs.Error("Invalid @Import format. Only accepts 1 or 2 parameters")
+				continue
+			}
+
+			p := parsedImport{}
+			p.importPath = iv[0]
+
+			if len(iv) == 2 {
+				p.importAlias = iv[1]
+			}
+
+			imports = append(imports, p)
+		}
+	}
+
+filterLoop:
+	for _, c := range lines {
+		t := strings.TrimSpace(strings.TrimLeft(c.Text, "//"))
+		if strings.HasPrefix(t, "@Filter") {
+			fv := getparams(strings.TrimSpace(strings.TrimLeft(t, "@Filter")))
+			if len(fv) < 3 {
+				logs.Error("Invalid @Filter format. Needs at least 3 parameters")
+				continue filterLoop
+			}
+
+			p := parsedFilter{}
+			p.pattern = fv[0]
+			posName := fv[1]
+			if pos, exists := routerHooks[posName]; exists {
+				p.pos = pos
+			} else {
+				logs.Error("Invalid @Filter pos: ", posName)
+				continue filterLoop
+			}
+
+			p.filter = fv[2]
+			fvParams := fv[3:]
+			for _, fvParam := range fvParams {
+				switch fvParam {
+				case "true":
+					p.params = append(p.params, true)
+				case "false":
+					p.params = append(p.params, false)
+				default:
+					logs.Error("Invalid @Filter param: ", fvParam)
+					continue filterLoop
+				}
+			}
+
+			filters = append(filters, p)
+		}
+	}
+
+	for _, c := range lines {
+		var pc = &parsedComment{}
+		pc.params = params
+		pc.filters = filters
+		pc.imports = imports
+
+		t := strings.TrimSpace(strings.TrimLeft(c.Text, "//"))
+		if strings.HasPrefix(t, "@router") {
+			t := strings.TrimSpace(strings.TrimLeft(c.Text, "//"))
+			matches := routeRegex.FindStringSubmatch(t)
+			if len(matches) == 3 {
+				pc.routerPath = matches[1]
+				methods := matches[2]
+				if methods == "" {
+					pc.methods = []string{"get"}
+					// pc.hasGet = true
+				} else {
+					pc.methods = strings.Split(methods, ",")
+					// pc.hasGet = strings.Contains(methods, "get")
+				}
+				pcs = append(pcs, pc)
+			} else {
+				return nil, errors.New("Router information is missing")
+			}
+		}
+	}
+	return
+}
+
+// direct copy from bee\g_docs.go
+// analysis params return []string
+// @Param	query		form	 string	true		"The email for login"
+// [query form string true "The email for login"]
+func getparams(str string) []string {
+	var s []rune
+	var j int
+	var start bool
+	var r []string
+	var quoted int8
+	for _, c := range str {
+		if unicode.IsSpace(c) && quoted == 0 {
+			if !start {
+				continue
+			} else {
+				start = false
+				j++
+				r = append(r, string(s))
+				s = make([]rune, 0)
+				continue
+			}
+		}
+
+		start = true
+		if c == '"' {
+			quoted ^= 1
+			continue
+		}
+		s = append(s, c)
+	}
+	if len(s) > 0 {
+		r = append(r, string(s))
+	}
+	return r
+}
+
+func genRouterCode(routersPath string) error {
+	logs.Info("generate router from comments")
+	var (
+		globalinfo   string
+		globalimport string
+		sortKey      []string
+	)
+	for k := range genInfoList {
+		sortKey = append(sortKey, k)
+	}
+	sort.Strings(sortKey)
+	for _, k := range sortKey {
+		cList := genInfoList[k]
+		sort.Sort(web.ControllerCommentsSlice(cList))
+		for _, c := range cList {
+			allmethod := "nil"
+			if len(c.AllowHTTPMethods) > 0 {
+				allmethod = "[]string{"
+				for _, m := range c.AllowHTTPMethods {
+					allmethod += `"` + m + `",`
+				}
+				allmethod = strings.TrimRight(allmethod, ",") + "}"
+			}
+
+			params := "nil"
+			if len(c.Params) > 0 {
+				params = "[]map[string]string{"
+				for _, p := range c.Params {
+					for k, v := range p {
+						params = params + `map[string]string{` + k + `:"` + v + `"},`
+					}
+				}
+				params = strings.TrimRight(params, ",") + "}"
+			}
+
+			methodParams := "param.Make("
+			if len(c.MethodParams) > 0 {
+				lines := make([]string, 0, len(c.MethodParams))
+				for _, m := range c.MethodParams {
+					lines = append(lines, fmt.Sprint(m))
+				}
+				methodParams += "\n				" +
+					strings.Join(lines, ",\n				") +
+					",\n			"
+			}
+			methodParams += ")"
+
+			imports := ""
+			if len(c.ImportComments) > 0 {
+				for _, i := range c.ImportComments {
+					var s string
+					if i.ImportAlias != "" {
+						s = fmt.Sprintf(`
+	%s "%s"`, i.ImportAlias, i.ImportPath)
+					} else {
+						s = fmt.Sprintf(`
+	"%s"`, i.ImportPath)
+					}
+					if !strings.Contains(globalimport, s) {
+						imports += s
+					}
+				}
+			}
+
+			filters := ""
+			if len(c.FilterComments) > 0 {
+				for _, f := range c.FilterComments {
+					filters += fmt.Sprintf(`                &beego.ControllerFilter{
+                    Pattern: "%s",
+                    Pos: %s,
+                    Filter: %s,
+                    ReturnOnOutput: %v,
+                    ResetParams: %v,
+                },`, f.Pattern, routerHooksMapping[f.Pos], f.Filter, f.ReturnOnOutput, f.ResetParams)
+				}
+			}
+
+			if filters == "" {
+				filters = "nil"
+			} else {
+				filters = fmt.Sprintf(`[]*beego.ControllerFilter{
+%s
+            }`, filters)
+			}
+
+			globalimport += imports
+
+			globalinfo = globalinfo + `
+    beego.GlobalControllerRouter["` + k + `"] = append(beego.GlobalControllerRouter["` + k + `"],
+        beego.ControllerComments{
+            Method: "` + strings.TrimSpace(c.Method) + `",
+            ` + "Router: `" + c.Router + "`" + `,
+            AllowHTTPMethods: ` + allmethod + `,
+            MethodParams: ` + methodParams + `,
+            Filters: ` + filters + `,
+            Params: ` + params + `})
+`
+		}
+	}
+
+	if globalinfo != "" {
+
+		routersPathDir := filepath.Dir(routersPath)
+		err := os.MkdirAll(routersPathDir, os.ModePerm)
+		if err != nil {
+			return err
+		}
+		f, err := os.Create(routersPath)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+
+		routersDir := RouterPkg.String()
+		if len(routersDir) == 0{
+			routersDir = "routers"
+		}
+
+		beeLogger.Log.Infof("using %s as routers file's package", routersDir)
+
+		content := strings.Replace(globalRouterTemplate, "{{.globalinfo}}", globalinfo, -1)
+		content = strings.Replace(content, "{{.routersDir}}", routersDir, -1)
+		content = strings.Replace(content, "{{.globalimport}}", globalimport, -1)
+		_, err = f.WriteString(content)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

go.mod

@@ -1,6 +1,6 @@
 module github.com/beego/bee/v2
 
-go 1.13
+go 1.16
 
 require (
 	github.com/beego/beego/v2 v2.0.1
@@ -14,6 +14,7 @@ require (
 	github.com/pelletier/go-toml v1.8.1
 	github.com/smartwalle/pongo2render v1.0.1
 	github.com/spf13/viper v1.7.0
+	golang.org/x/tools v0.0.0-20201211185031-d93e913c1a58
 	gopkg.in/yaml.v2 v2.3.0
 )