Merge pull request #70 from hectorj/process-pkg

Allow to process entire package at once

easyjson/main.go

@@ -1,9 +1,11 @@
 package main
 
 import (
+	"errors"
 	"flag"
 	"fmt"
 	"os"
+	"path/filepath"
 	"strings"
 
 	"github.com/mailru/easyjson/bootstrap"
@@ -23,18 +25,28 @@ var leaveTemps = flag.Bool("leave_temps", false, "do not delete temporary files"
 var stubs = flag.Bool("stubs", false, "only generate stubs for marshallers/unmarshallers methods")
 var noformat = flag.Bool("noformat", false, "do not run 'gofmt -w' on output file")
 var specifiedName = flag.String("output_filename", "", "specify the filename of the output")
+var processPkg = flag.Bool("pkg", false, "process the whole package instead of just the given file")
 
 func generate(fname string) (err error) {
+	fInfo, err := os.Stat(fname)
+	if err != nil {
+		return err
+	}
+
 	p := parser.Parser{AllStructs: *allStructs}
-	if err := p.Parse(fname); err != nil {
+	if err := p.Parse(fname, fInfo.IsDir()); err != nil {
 		return fmt.Errorf("Error parsing %v: %v", fname, err)
 	}
 
 	var outName string
-	if s := strings.TrimSuffix(fname, ".go"); s == fname {
-		return fmt.Errorf("Filename must end in '.go'")
+	if fInfo.IsDir() {
+		outName = filepath.Join(fname, p.PkgName+"_easyjson.go")
 	} else {
-		outName = s + "_easyjson.go"
+		if s := strings.TrimSuffix(fname, ".go"); s == fname {
+			return errors.New("Filename must end in '.go'")
+		} else {
+			outName = s + "_easyjson.go"
+		}
 	}
 
 	if *specifiedName != "" {
@@ -67,6 +79,10 @@ func main() {
 	files := flag.Args()
 
 	gofile := os.Getenv("GOFILE")
+	if *processPkg {
+		gofile = filepath.Dir(gofile)
+	}
+
 	if len(files) == 0 && gofile != "" {
 		files = []string{gofile}
 	} else if len(files) == 0 {

parser/parser.go

@@ -34,6 +34,8 @@ func (p *Parser) needType(comments string) bool {
 
 func (v *visitor) Visit(n ast.Node) (w ast.Visitor) {
 	switch n := n.(type) {
+	case *ast.Package:
+		return v
 	case *ast.File:
 		v.PkgName = n.Name.String()
 		return v
@@ -61,18 +63,29 @@ func (v *visitor) Visit(n ast.Node) (w ast.Visitor) {
 	return nil
 }
 
-func (p *Parser) Parse(fname string) error {
+func (p *Parser) Parse(fname string, isDir bool) error {
 	var err error
-	if p.PkgPath, err = getPkgPath(fname); err != nil {
+	if p.PkgPath, err = getPkgPath(fname, isDir); err != nil {
 		return err
 	}
 
 	fset := token.NewFileSet()
-	f, err := parser.ParseFile(fset, fname, nil, parser.ParseComments)
-	if err != nil {
-		return err
-	}
+	if isDir {
+		packages, err := parser.ParseDir(fset, fname, nil, parser.ParseComments)
+		if err != nil {
+			return err
+		}
 
-	ast.Walk(&visitor{Parser: p}, f)
+		for _, pckg := range packages {
+			ast.Walk(&visitor{Parser: p}, pckg)
+		}
+	} else {
+		f, err := parser.ParseFile(fset, fname, nil, parser.ParseComments)
+		if err != nil {
+			return err
+		}
+
+		ast.Walk(&visitor{Parser: p}, f)
+	}
 	return nil
 }

parser/parser_unix.go

@@ -9,7 +9,7 @@ import (
 	"strings"
 )
 
-func getPkgPath(fname string) (string, error) {
+func getPkgPath(fname string, isDir bool) (string, error) {
 	if !path.IsAbs(fname) {
 		pwd, err := os.Getwd()
 		if err != nil {
@@ -21,7 +21,11 @@ func getPkgPath(fname string) (string, error) {
 	for _, p := range strings.Split(os.Getenv("GOPATH"), ":") {
 		prefix := path.Join(p, "src") + "/"
 		if rel := strings.TrimPrefix(fname, prefix); rel != fname {
-			return path.Dir(rel), nil
+			if !isDir {
+				return path.Dir(rel), nil
+			} else {
+				return path.Clean(rel), nil
+			}
 		}
 	}