mm-ast-refactor-talk.go

package main

import (
	"bytes"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"io/ioutil"
	"strings"
)

func fix(dir string) {
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, dir, nil, parser.ParseComments)
	if err != nil {
		panic(err)
	}

	for _, pkg := range pkgs {
		for fileName, file := range pkg.Files {
			ast.Inspect(file, func(node ast.Node) bool {
				var list *[]ast.Stmt
				switch x := node.(type) {
				case *ast.BlockStmt:
					list = &x.List
				case *ast.CaseClause:
					list = &x.Body
				case *ast.CommClause:
					list = &x.Body
				}
				if list == nil {
					return true
				}
				var newList []ast.Stmt
				var newAssignment *ast.AssignStmt
				var newFunction ast.Expr
				for _, stmt := range *list {
					if assignment, ok := stmt.(*ast.AssignStmt); ok && newAssignment != nil {
						if star, ok := assignment.Lhs[0].(*ast.StarExpr); ok {
							fset.File(assignment.TokPos).MergeLine(fset.Position(newAssignment.TokPos).Line)
							*newAssignment = ast.AssignStmt{
								Lhs: []ast.Expr{star.X},
								Rhs: []ast.Expr{
									&ast.CallExpr{
										Fun:  newFunction,
										Args: assignment.Rhs,
									},
								},
								Tok: assignment.Tok,
							}
							continue
						}
					}

					newList = append(newList, stmt)
					newAssignment = nil
					if assignment, ok := stmt.(*ast.AssignStmt); ok && len(assignment.Rhs) == 1 {
						if call, ok := assignment.Rhs[0].(*ast.CallExpr); ok && len(call.Args) == 1 {
							if ident, ok := call.Fun.(*ast.Ident); ok && ident.Name == "new" {
								if typeIdent, ok := call.Args[0].(*ast.Ident); ok {
									switch typeIdent.Name {
									case "string", "int", "int64", "bool":
										ident := &ast.Ident{
											Name: "New" + strings.Title(typeIdent.Name),
										}
										newFunction = ident
										if dir != "./model" {
											newFunction = &ast.SelectorExpr{
												X:   &ast.Ident{Name: "model"},
												Sel: ident,
											}
										}
										newAssignment = assignment
									}
								}
							}
						}
					}
				}
				*list = newList
				return true
			})

			buf := &bytes.Buffer{}
			format.Node(buf, fset, file)
			if err := ioutil.WriteFile(fileName, buf.Bytes(), 0664); err != nil {
				panic(err)
			}
		}
	}
}

func main() {
	for _, dir := range []string{"./api", "./api4", "./app", "./model"} {
		fix(dir)
	}
}