Golang astutil.AddImport函数代码示例

// addCode searches for main func in data, and updates AST code
// adding tracing functions.
func addCode(path string) ([]byte, error) {
	var conf loader.Config
	if _, err := conf.FromArgs([]string{path}, false); err != nil {
		return nil, err
	}

	prog, err := conf.Load()
	if err != nil {
		return nil, err
	}

	// check if runtime/trace already imported
	for i, _ := range prog.Imported {
		if i == "runtime/trace" {
			return nil, ErrImported
		}
	}

	pkg := prog.Created[0]

	// TODO: find file with main func inside
	astFile := pkg.Files[0]

	// add imports
	astutil.AddImport(prog.Fset, astFile, "os")
	astutil.AddImport(prog.Fset, astFile, "runtime/trace")
	astutil.AddImport(prog.Fset, astFile, "time")

	// add start/stop code
	ast.Inspect(astFile, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.FuncDecl:
			// find 'main' function
			if x.Name.Name == "main" && x.Recv == nil {
				stmts := createTraceStmts()
				stmts = append(stmts, x.Body.List...)
				x.Body.List = stmts
				return true
			}
		}
		return true
	})

	var buf bytes.Buffer
	err = printer.Fprint(&buf, prog.Fset, astFile)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func args(filename string) {
	b, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}

	p := parser.ParseFromString(filename, string(b)+"\n")

	a := generator.GenerateAST(p)

	fset := token.NewFileSet()

	defaultImports := []string{"github.com/gsp-lang/stdlib/prelude", "github.com/gsp-lang/gsp/core"}
	for _, defaultImport := range defaultImports {
		split := strings.Split(defaultImport, "/")
		pkgName := split[len(split)-1]
		if !(a.Name.Name == "prelude" && pkgName == "prelude") {
			astutil.AddImport(fset, a, defaultImport)
		}
	}

	if a.Name.Name != "prelude" {
		a.Decls = append(a.Decls, &ast.GenDecl{
			Tok: token.VAR,
			Specs: []ast.Spec{&ast.ValueSpec{
				Names:  []*ast.Ident{&ast.Ident{Name: "_"}},
				Values: []ast.Expr{&ast.Ident{Name: "prelude.Len"}},
			}},
		})
	}

	var buf bytes.Buffer
	printer.Fprint(&buf, fset, a)
	fmt.Printf("%s\n", buf.String())
}

func args(filename string) {
	b, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}

	p := parser.ParseFromString(filename, string(b)+"\n")

	a := generator.GenerateAST(p)

	fset := token.NewFileSet()

	defaultImports := []string{"github.com/gsp-lang/stdlib/prelude", "github.com/gsp-lang/gsp/core"}
	for _, defaultImport := range defaultImports {
		split := strings.Split(defaultImport, "/")
		pkgName := split[len(split)-1]
		if !(a.Name.Name == "prelude" && pkgName == "prelude") {
			if pkgName == "prelude" {
				astutil.AddNamedImport(fset, a, "_", defaultImport)
			} else {
				astutil.AddImport(fset, a, defaultImport)
			}
		}
	}

	var buf bytes.Buffer
	printer.Fprint(&buf, fset, a)
	fmt.Printf("%s\n", buf.String())
}

// Inline replaces each instance of identifier k with v.Ident in ast.File f,
// for k, v := range m.
// For all inlines that were triggeres it also adds imports from v.Imports to f.
// In addition, it removes top level type declarations of the form
// type k ...
// for all k in m.
//
// Every k in m should be a valid identifier.
// Every v.Ident should be a valid expression.
func Inline(fset *token.FileSet, f *ast.File, m map[string]Target) error {
	// Build the inline map.
	im := map[string]reflect.Value{}
	for k, v := range m {
		expr, err := parser.ParseExpr(k)
		if err != nil {
			return fmt.Errorf("failed to parse `%s`: %s", k, err)
		}
		if _, ok := expr.(*ast.Ident); !ok {
			return fmt.Errorf("expected identifier, got %s which is %T", k, expr)
		}
		expr, err = parser.ParseExpr(v.Ident)
		if err != nil {
			return fmt.Errorf("failed to parse `%s`: %s", v.Ident, err)
		}
		s := v.Ident
		if _, ok := expr.(*ast.StarExpr); ok {
			s = fmt.Sprintf("(%s)", s)
		}
		im[k] = reflect.ValueOf(ast.Ident{Name: s})
	}
	// Filter `type XXX ...` declarations out if we are inlining XXX.
	cmap := ast.NewCommentMap(fset, f, f.Comments)
	to := 0
	for _, d := range f.Decls {
		skip := false
		if t, ok := d.(*ast.GenDecl); ok {
			for _, s := range t.Specs {
				ts, ok := s.(*ast.TypeSpec)
				if !ok {
					continue
				}
				if _, ok = im[ts.Name.String()]; ok {
					skip = true
				}
			}
		}
		if !skip {
			f.Decls[to] = d
			to++
		}
	}
	if to != len(f.Decls) {
		f.Decls = f.Decls[:to]
		// Remove comments for the declarations that were filtered out.
		f.Comments = cmap.Filter(f).Comments()
	}
	// Add imports for the inlines that were triggered.
	for k := range inline(im, f) {
		for _, imp := range m[k].Imports {
			astutil.AddImport(fset, f, imp)
		}
	}
	return nil
}

func actionImport(s *Session, arg string) error {
	if arg == "" {
		return fmt.Errorf("arg required")
	}

	path := strings.Trim(arg, `"`)

	// check if the package specified by path is importable
	_, err := s.Types.Importer.Import(path)
	if err != nil {
		return err
	}

	astutil.AddImport(s.Fset, s.File, path)

	return nil
}

func addImports(file *ast.File, fset *token.FileSet, dirPath string) (*ast.File, *token.FileSet, error) {
	imports, err := getImports(dirPath, fset)
	if err != nil {
		return nil, nil, err
	}

	for _, s := range imports {
		unquotedPath, err := strconv.Unquote(s.Path.Value)
		if err != nil {
			return nil, nil, err
		}

		if s.Name != nil {
			astutil.AddNamedImport(fset, file, s.Name.Name, unquotedPath)
			continue
		}

		astutil.AddImport(fset, file, unquotedPath)
	}

	return file, fset, nil
}

//.........这里部分代码省略.........

	abs, err := filepath.Abs(filename)
	if err != nil {
		return nil, err
	}
	srcDir := path.Dir(abs)

	// collect potential uses of packages.
	var visitor visitFn
	visitor = visitFn(func(node ast.Node) ast.Visitor {
		if node == nil {
			return visitor
		}
		switch v := node.(type) {
		case *ast.ImportSpec:
			if v.Name != nil {
				decls[v.Name.Name] = v
			} else {
				local := importPathToName(strings.Trim(v.Path.Value, `\"`), srcDir)
				decls[local] = v
			}
		case *ast.SelectorExpr:
			xident, ok := v.X.(*ast.Ident)
			if !ok {
				break
			}
			if xident.Obj != nil {
				// if the parser can resolve it, it's not a package ref
				break
			}
			pkgName := xident.Name
			if refs[pkgName] == nil {
				refs[pkgName] = make(map[string]bool)
			}
			if decls[pkgName] == nil {
				refs[pkgName][v.Sel.Name] = true
			}
		}
		return visitor
	})
	ast.Walk(visitor, f)

	// Nil out any unused ImportSpecs, to be removed in following passes
	unusedImport := map[string]string{}
	for pkg, is := range decls {
		if refs[pkg] == nil && pkg != "_" && pkg != "." {
			name := ""
			if is.Name != nil {
				name = is.Name.Name
			}
			unusedImport[strings.Trim(is.Path.Value, `"`)] = name
		}
	}
	for ipath, name := range unusedImport {
		if ipath == "C" {
			// Don't remove cgo stuff.
			continue
		}
		astutil.DeleteNamedImport(fset, f, name, ipath)
	}

	// Search for imports matching potential package references.
	searches := 0
	type result struct {
		ipath string
		name  string
		err   error
	}
	results := make(chan result)
	for pkgName, symbols := range refs {
		if len(symbols) == 0 {
			continue // skip over packages already imported
		}
		go func(pkgName string, symbols map[string]bool) {
			ipath, rename, err := findImport(pkgName, symbols, filename)
			r := result{ipath: ipath, err: err}
			if rename {
				r.name = pkgName
			}
			results <- r
		}(pkgName, symbols)
		searches++
	}
	for i := 0; i < searches; i++ {
		result := <-results
		if result.err != nil {
			return nil, result.err
		}
		if result.ipath != "" {
			if result.name != "" {
				astutil.AddNamedImport(fset, f, result.name, result.ipath)
			} else {
				astutil.AddImport(fset, f, result.ipath)
			}
			added = append(added, result.ipath)
		}
	}

	return added, nil
}

// Transform applies the transformation to the specified parsed file,
// whose type information is supplied in info, and returns the number
// of replacements that were made.
//
// It mutates the AST in place (the identity of the root node is
// unchanged), and may add nodes for which no type information is
// available in info.
//
// Derived from rewriteFile in $GOROOT/src/cmd/gofmt/rewrite.go.
//
func (tr *Transformer) Transform(info *types.Info, pkg *types.Package, file *ast.File) int {
	if !tr.seenInfos[info] {
		tr.seenInfos[info] = true
		mergeTypeInfo(tr.info, info)
	}
	tr.currentPkg = pkg
	tr.nsubsts = 0

	if tr.verbose {
		fmt.Fprintf(os.Stderr, "before: %s\n", astString(tr.fset, tr.before))
		fmt.Fprintf(os.Stderr, "after: %s\n", astString(tr.fset, tr.after))
	}

	var f func(rv reflect.Value) reflect.Value
	f = func(rv reflect.Value) reflect.Value {
		// don't bother if val is invalid to start with
		if !rv.IsValid() {
			return reflect.Value{}
		}

		rv = apply(f, rv)

		e := rvToExpr(rv)
		if e != nil {
			savedEnv := tr.env
			tr.env = make(map[string]ast.Expr) // inefficient!  Use a slice of k/v pairs

			if tr.matchExpr(tr.before, e) {
				if tr.verbose {
					fmt.Fprintf(os.Stderr, "%s matches %s",
						astString(tr.fset, tr.before), astString(tr.fset, e))
					if len(tr.env) > 0 {
						fmt.Fprintf(os.Stderr, " with:")
						for name, ast := range tr.env {
							fmt.Fprintf(os.Stderr, " %s->%s",
								name, astString(tr.fset, ast))
						}
					}
					fmt.Fprintf(os.Stderr, "\n")
				}
				tr.nsubsts++

				// Clone the replacement tree, performing parameter substitution.
				// We update all positions to n.Pos() to aid comment placement.
				rv = tr.subst(tr.env, reflect.ValueOf(tr.after),
					reflect.ValueOf(e.Pos()))
			}
			tr.env = savedEnv
		}

		return rv
	}
	file2 := apply(f, reflect.ValueOf(file)).Interface().(*ast.File)

	// By construction, the root node is unchanged.
	if file != file2 {
		panic("BUG")
	}

	// Add any necessary imports.
	// TODO(adonovan): remove no-longer needed imports too.
	if tr.nsubsts > 0 {
		pkgs := make(map[string]*types.Package)
		for obj := range tr.importedObjs {
			pkgs[obj.Pkg().Path()] = obj.Pkg()
		}

		for _, imp := range file.Imports {
			path, _ := strconv.Unquote(imp.Path.Value)
			delete(pkgs, path)
		}
		delete(pkgs, pkg.Path()) // don't import self

		// NB: AddImport may completely replace the AST!
		// It thus renders info and tr.info no longer relevant to file.
		var paths []string
		for path := range pkgs {
			paths = append(paths, path)
		}
		sort.Strings(paths)
		for _, path := range paths {
			astutil.AddImport(tr.fset, file, path)
		}
	}

	tr.currentPkg = nil

	return tr.nsubsts
}

// ProcessFileAST processes the files using golang's AST parser
func ProcessFileAST(filePath string, from string, to string) {

	//Colors to be used on the console
	red := ansi.ColorCode("red+bh")
	white := ansi.ColorCode("white+bh")
	yellow := ansi.ColorCode("yellow+bh")
	blackOnWhite := ansi.ColorCode("black+b:white+h")
	//Reset the color
	reset := ansi.ColorCode("reset")

	fmt.Println(blackOnWhite+"Processing file", filePath, "in SAFE MODE", reset)

	// New FileSet to parse the go file to
	fSet := token.NewFileSet()

	// Parse the file
	file, err := parser.ParseFile(fSet, filePath, nil, 0)
	if err != nil {
		fmt.Println(err)
	}

	// Get the list of imports from the ast
	imports := astutil.Imports(fSet, file)

	// Keep track of number of changes
	numChanges := 0

	// Iterate through the imports array
	for _, mPackage := range imports {
		for _, mImport := range mPackage {
			// Since astutil returns the path string with quotes, remove those
			importString := strings.TrimSuffix(strings.TrimPrefix(mImport.Path.Value, "\""), "\"")

			// If the path matches the oldpath, replace it with the new one
			if strings.Contains(importString, from) {
				//If it needs to be replaced, increase numChanges so we can write the file later
				numChanges++

				// Join the path of the import package with the remainder from the old one after removing the old import package
				replacePackage := strings.Replace(importString, from, to, -1)

				fmt.Println(red +
					"Updating import " +
					reset + white +
					importString +
					reset + red +
					" to " +
					reset + white +
					replacePackage +
					reset)

				// Remove the old import and replace it with the replacement
				astutil.DeleteImport(fSet, file, importString)
				astutil.AddImport(fSet, file, replacePackage)
			}
		}
	}

	// If the number of changes are more than 0, write file
	if numChanges > 0 {
		// Print the new AST tree to a new output buffer
		var outputBuffer bytes.Buffer
		printer.Fprint(&outputBuffer, fSet, file)

		ioutil.WriteFile(filePath, outputBuffer.Bytes(), os.ModePerm)
		fmt.Println(yellow+
			"File",
			filePath,
			"saved after",
			numChanges,
			"changes",
			reset, "\n\n")
	} else {
		fmt.Println(yellow+
			"No changes to write on this file.",
			reset, "\n\n")
	}
}

func (d *DirectiveList) Save() error {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, d.file, nil, 0)
	if err != nil {
		return ErrInvalidDirectiveFile
	}

	imps := astutil.Imports(fset, f)
	for _, imp := range imps[0] {
		name := ""
		if imp.Name != nil && imp.Name.String() != "." {
			name = imp.Name.String()
		}
		path, _ := strconv.Unquote(imp.Path.Value)

		if name != "" {
			astutil.DeleteNamedImport(fset, f, name, path)
		} else {
			astutil.DeleteImport(fset, f, path)
		}
	}

	astutil.AddImport(fset, f, "github.com/mholt/caddy/caddy/https")
	astutil.AddImport(fset, f, "github.com/mholt/caddy/caddy/parse")
	astutil.AddImport(fset, f, "github.com/mholt/caddy/caddy/setup")
	astutil.AddImport(fset, f, "github.com/mholt/caddy/middleware")

	importName := ""
	for _, dir := range d.list {
		if dir.Removed {
			continue
		}
		if dir.Core == false && len(dir.ImportPath) != 0 {
			importName = strings.Replace(dir.Name, "-", "_", -1)
			astutil.AddImport(fset, f, fmt.Sprintf("{{import-%s}}", importName))
		}
	}

	var buf bytes.Buffer
	err = printer.Fprint(&buf, fset, f)
	if err != nil {
		return err
	}

	out := buf.String()

	f, err = parser.ParseFile(token.NewFileSet(), "", out, 0)
	node, ok := f.Scope.Lookup("directiveOrder").Decl.(ast.Node)
	if !ok {
		return ErrInvalidDirectiveFile
	}

	begin := out[0 : node.Pos()-1]
	end := out[node.End():len(out)]

	dirOrder := ""
	for _, dir := range d.list {
		if dir.Removed {
			continue
		}

		comment := ""
		if dir.Active == false {
			comment = "//@caddyext "
		}

		importName = strings.Replace(dir.Name, "-", "_", -1)

		begin = strings.Replace(begin, fmt.Sprintf(`"{{import-%s}}"`, importName), fmt.Sprintf(`%s%s "%s"`, comment, importName, dir.ImportPath), -1)

		if dir.Core == true {
			dirOrder = dirOrder + fmt.Sprintf(`	%s{"%s", %s},`+"\n", comment, dir.Name, dir.Setup)
		} else {
			dirOrder = dirOrder + fmt.Sprintf(`	%s{"%s", %s.Setup},`+"\n", comment, dir.Name, importName)
		}
	}

	out = begin + dirOrderStart + dirOrder + dirOrderEnd + end

	return ioutil.WriteFile(d.file, []byte(out), os.FileMode(0660))
}

// gen is code generation function that insert custom directives at runtime.
func gen(middlewares features.Middlewares) custombuild.CodeGenFunc {
	return func(src string, packages []string) (err error) {
		// prevent possible panic from assertions.
		defer func() {
			if recover() != nil {
				err = errParse
			}
		}()

		// if no middleware is added, no code generation needed.
		if len(middlewares) == 0 {
			return nil
		}

		fset := token.NewFileSet()
		file := filepath.Join(src, directivesFile)
		f, err := parser.ParseFile(fset, file, nil, 0)
		if err != nil {
			return err
		}
		packageNames, err := getPackageNames(middlewares.Packages())
		if err != nil {
			return err
		}
		for _, m := range middlewares {
			astutil.AddImport(fset, f, m.Package)
		}
		var buf bytes.Buffer
		err = printer.Fprint(&buf, fset, f)
		if err != nil {
			return err
		}

		out := buf.String()

		for _, mid := range middlewares {
			f, err = parser.ParseFile(token.NewFileSet(), "", out, 0)
			node, ok := f.Scope.Lookup("directiveOrder").Decl.(ast.Node)
			if !ok {
				return errParse
			}

			snippet := fmt.Sprintf(`{"%s", %s.Setup},`+"\n", mid.Directive, packageNames[mid.Package])

			// add to end of directives.
			end := int(node.End()) - 2

			// if after is set, locate directive and add after it.
			after := getPrevDirective(mid.Directive)
			if after != "" {
				found := false
				c := node.(*ast.ValueSpec).Values[0].(*ast.CompositeLit)
				for _, m := range c.Elts {
					directive := m.(*ast.CompositeLit).Elts[0].(*ast.BasicLit)
					if strconv.Quote(after) == directive.Value && !found {
						end = int(m.End()) + 1
						found = true
					}
					// check if directive exists
					if strconv.Quote(mid.Directive) == directive.Value {
						return fmt.Errorf("Directive '%s' exists in Caddy core, use a distinct name.", mid.Directive)
					}
				}
				if !found {
					return fmt.Errorf("Cannot place afer %s, directive '%s' not found.", config.After, config.After)
				}
			}

			out = out[:end] + snippet + out[end:]
		}

		return ioutil.WriteFile(file, []byte(out), os.FileMode(0660))
	}
}