parse.go

package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
)

var (
	gopath string
	goroot string
)

func init() {
	gopath = os.Getenv("GOPATH")
	if gopath == "" {
		fmt.Fprintf(os.Stderr, "ERROR: GOPATH environment not found\n")
		os.Exit(2)
	}
	goroot = os.Getenv("GOROOT")
	if goroot == "" {
		fmt.Fprintf(os.Stderr, "ERROR: GOROOT environment not found\n")
		os.Exit(2)
	}
}

func noTestsFilter(fi os.FileInfo) bool {
	name := fi.Name()
	if strings.HasSuffix(name, "_test.go") {
		return false
	}
	return strings.HasSuffix(name, ".go")
}

// convert a package name to a filesystem path
func packagePath(pkg string) string {
	pkg = filepath.Join(strings.Split(pkg, "/")...)
	if strings.Contains(pkg, ".") {
		return filepath.Join(gopath, "src", pkg)
	}
	return filepath.Join(goroot, "src", pkg)
}

// parse a package but parse it's dependencies first
func parsePackage(pkg string) (*token.FileSet, map[string]*ast.Package, error) {
	fset := token.NewFileSet() // the fset will accumulate dependency parses
	pkgSeen := map[string]bool{}
	var parseImport func(string) (map[string]*ast.Package, error)
	// recursively parse all imports (depth-first)
	parseImport = func(pkg string) (map[string]*ast.Package, error) {
		if pkgSeen[pkg] || pkg == "C" {
			return nil, nil
		}
		pkgSeen[pkg] = true
		pkgDir := packagePath(pkg)
		// get the package's imports first. 
		pkgs, err := parser.ParseDir(fset, pkgDir, noTestsFilter, parser.ImportsOnly)
		if err != nil {
			return nil, err
		}
		imports := map[string]bool{}
		for _, p := range pkgs {
			for _, file := range p.Files {
				for i := range file.Imports {
					name := file.Imports[i].Path.Value
					name = strings.Trim(name, `"`)
					if imports[name] {
					    continue // we've already seen this import within this package
					}
					imports[name] = true
					// parse this import
					if _, err := parseImport(name); err != nil {
						return nil, err
					}
				}
			}
		}
		return parser.ParseDir(fset, pkgDir, noTestsFilter, 0)
	}
	astPkg, err := parseImport(pkg)
	return fset, astPkg, err
}

type indentVisitor struct {
	fset *token.FileSet
}

// print all ast.CallExpr
func (v *indentVisitor) Visit(n ast.Node) ast.Visitor {
	switch n := n.(type) {
	case *ast.CallExpr:
		ast.Print(v.fset, n)
	}
	return v
}

func main() {
	fset, pkg, err := parsePackage("github.com/ericchiang/pup")
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(2)
	}
	if len(pkg) != 1 {
		formatStr := "returned a bad number of packages (expected 1): %d"
		fmt.Fprintf(os.Stderr, formatStr, len(pkg))
		os.Exit(2)
	}
	v := &indentVisitor{fset}
	for _, pkgData := range pkg {
		ast.Walk(v, pkgData)
	}
}