gopfmt: support classfile

cmd/gopfmt/gopfmt.go

@@ -31,19 +31,14 @@ import (
 
 	goformat "go/format"
 
+	"github.com/goplus/gop"
 	"github.com/goplus/gop/format"
 )
 
 var (
 	procCnt    = 0
 	walkSubDir = false
-	extGops    = map[string]struct{}{
-		".go":  {},
-		".gop": {},
-		".spx": {},
-		".gmx": {},
-	}
-	rootDir = ""
+	rootDir    = ""
 )
 
 var (
@@ -61,7 +56,7 @@ func report(err error) {
 	os.Exit(2)
 }
 
-func processFile(filename string, in io.Reader, out io.Writer) error {
+func processFile(filename string, class bool, in io.Reader, out io.Writer) error {
 	if in == nil {
 		var err error
 		in, err = os.Open(filename)
@@ -87,7 +82,7 @@ func processFile(filename string, in io.Reader, out io.Writer) error {
 		}
 		res = buf.Bytes()
 	} else {
-		res, err = format.Source(src, filename)
+		res, err = format.Source(src, class, filename)
 	}
 	if err != nil {
 		return err
@@ -118,7 +113,15 @@ func processFile(filename string, in io.Reader, out io.Writer) error {
 	return err
 }
 
-func walk(path string, d fs.DirEntry, err error) error {
+type walker struct {
+	dirMap map[string]func(ext string) (ok bool, class bool)
+}
+
+func newWalker() *walker {
+	return &walker{dirMap: make(map[string]func(ext string) (ok bool, class bool))}
+}
+
+func (w *walker) walk(path string, d fs.DirEntry, err error) error {
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
 	} else if d.IsDir() {
@@ -127,10 +130,41 @@ func walk(path string, d fs.DirEntry, err error) error {
 		}
 	} else {
 		// Directories are walked, ignoring non-Gop files.
+		dir, _ := filepath.Split(path)
+		fn, ok := w.dirMap[dir]
+		if !ok {
+			if mod, err := gop.LoadMod(path, nil, &gop.Config{DontUpdateGoMod: true}); err == nil {
+				fn = func(ext string) (ok bool, class bool) {
+					switch ext {
+					case ".go", ".gop":
+						ok = true
+					case ".gopx", ".spx", ".gmx":
+						ok, class = true, true
+					default:
+						_, class = mod.IsClass(ext)
+						if class {
+							ok = true
+						}
+					}
+					return
+				}
+			} else {
+				fn = func(ext string) (ok bool, class bool) {
+					switch ext {
+					case ".go", ".gop":
+						ok = true
+					case ".gopx", ".spx", ".gmx":
+						ok, class = true, true
+					}
+					return
+				}
+			}
+			w.dirMap[dir] = fn
+		}
 		ext := filepath.Ext(path)
-		if _, ok := extGops[ext]; ok {
+		if ok, class := fn(ext); ok {
 			procCnt++
-			if err = processFile(path, nil, os.Stdout); err != nil {
+			if err = processFile(path, class, nil, os.Stdout); err != nil {
 				report(err)
 			}
 		}
@@ -148,20 +182,20 @@ func main() {
 			report(fmt.Errorf("error: cannot use -w with standard input"))
 			return
 		}
-		if err := processFile("<standard input>", os.Stdin, os.Stdout); err != nil {
+		if err := processFile("<standard input>", false, os.Stdin, os.Stdout); err != nil {
 			report(err)
 		}
 		return
 	}
-
+	walker := newWalker()
 	for _, path := range args {
 		walkSubDir = strings.HasSuffix(path, "/...")
 		if walkSubDir {
 			path = path[:len(path)-4]
 		}
 		procCnt = 0
 		rootDir = path
-		filepath.WalkDir(path, walk)
+		filepath.WalkDir(path, walker.walk)
 		if procCnt == 0 {
 			fmt.Println("no Go+ files in", path)
 		}

cmd/internal/gopfmt/fmt.go

@@ -27,6 +27,8 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/goplus/gop"
+
 	"github.com/goplus/gop/cmd/internal/base"
 	"github.com/goplus/gop/format"
 
@@ -57,16 +59,10 @@ func init() {
 var (
 	procCnt    = 0
 	walkSubDir = false
-	extGops    = map[string]struct{}{
-		".go":  {},
-		".gop": {},
-		".spx": {},
-		".gmx": {},
-	}
-	rootDir = ""
+	rootDir    = ""
 )
 
-func gopfmt(path string, smart, mvgo bool) (err error) {
+func gopfmt(path string, class, smart, mvgo bool) (err error) {
 	src, err := ioutil.ReadFile(path)
 	if err != nil {
 		return
@@ -88,7 +84,7 @@ func gopfmt(path string, smart, mvgo bool) (err error) {
 			}
 			target = buf.Bytes()
 		} else {
-			target, err = format.Source(src, path)
+			target, err = format.Source(src, class, path)
 		}
 	}
 	if err != nil {
@@ -127,23 +123,62 @@ func writeFileWithBackup(path string, target []byte) (err error) {
 	return os.Rename(tmpfile, path)
 }
 
-func walk(path string, d fs.DirEntry, err error) error {
+type walker struct {
+	dirMap map[string]func(ext string) (ok, class bool)
+}
+
+func newWalker() *walker {
+	return &walker{dirMap: make(map[string]func(ext string) (ok, class bool))}
+}
+
+func (w *walker) walk(path string, d fs.DirEntry, err error) error {
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
 	} else if d.IsDir() {
 		if !walkSubDir && path != rootDir {
 			return filepath.SkipDir
 		}
 	} else {
+		dir, _ := filepath.Split(path)
+		fn, ok := w.dirMap[dir]
+		if !ok {
+			if mod, err := gop.LoadMod(path, nil, &gop.Config{DontUpdateGoMod: true}); err == nil {
+				fn = func(ext string) (ok bool, class bool) {
+					switch ext {
+					case ".go", ".gop":
+						ok = true
+					case ".gopx", ".spx", ".gmx":
+						ok, class = true, true
+					default:
+						_, class = mod.IsClass(ext)
+						if class {
+							ok = true
+						}
+					}
+					return
+				}
+			} else {
+				fn = func(ext string) (ok bool, class bool) {
+					switch ext {
+					case ".go", ".gop":
+						ok = true
+					case ".gopx", ".spx", ".gmx":
+						ok, class = true, true
+					}
+					return
+				}
+			}
+			w.dirMap[dir] = fn
+		}
 		ext := filepath.Ext(path)
 		smart := *flagSmart
 		mvgo := smart && *flagMoveGo
-		if _, ok := extGops[ext]; ok && (!mvgo || ext == ".go") {
+		if ok, class := fn(ext); ok && (!mvgo || ext == ".go") {
 			procCnt++
 			if *flagNotExec {
 				fmt.Println("gop fmt", path)
 			} else {
-				err = gopfmt(path, smart && (mvgo || ext != ".go"), mvgo)
+				err = gopfmt(path, class, smart && (mvgo || ext != ".go"), mvgo)
 				if err != nil {
 					report(err)
 				}
@@ -167,6 +202,7 @@ func runCmd(cmd *base.Command, args []string) {
 	if narg < 1 {
 		cmd.Usage(os.Stderr)
 	}
+	walker := newWalker()
 	for i := 0; i < narg; i++ {
 		path := flag.Arg(i)
 		walkSubDir = strings.HasSuffix(path, "/...")
@@ -175,7 +211,7 @@ func runCmd(cmd *base.Command, args []string) {
 		}
 		procCnt = 0
 		rootDir = path
-		filepath.WalkDir(path, walk)
+		filepath.WalkDir(path, walker.walk)
 		if procCnt == 0 {
 			fmt.Println("no Go+ files in", path)
 		}

format/format.go

@@ -50,7 +50,6 @@ const parserMode = parser.ParseComments
 //
 // The function may return early (before the entire result is written)
 // and return a formatting error, for instance due to an incorrect AST.
-//
 func Node(dst io.Writer, fset *token.FileSet, node interface{}) error {
 	// Determine if we have a complete source file (file != nil).
 	var file *ast.File
@@ -101,14 +100,13 @@ func Node(dst io.Writer, fset *token.FileSet, node interface{}) error {
 // is applied to the result (such that it has the same leading and trailing
 // space as src), and the result is indented by the same amount as the first
 // line of src containing code. Imports are not sorted for partial source files.
-//
-func Source(src []byte, filename ...string) ([]byte, error) {
+func Source(src []byte, class bool, filename ...string) ([]byte, error) {
 	var fname string
 	if filename != nil {
 		fname = filename[0]
 	}
 	fset := token.NewFileSet()
-	file, sourceAdj, indentAdj, err := parse(fset, fname, src, true)
+	file, sourceAdj, indentAdj, err := parse(fset, fname, src, class, true)
 	if err != nil {
 		return nil, err
 	}

format/internal.go

@@ -32,13 +32,17 @@ import (
 
 // parse parses src, which was read from the named file,
 // as a Go source file, declaration, or statement list.
-func parse(fset *token.FileSet, filename string, src []byte, fragmentOk bool) (
+func parse(fset *token.FileSet, filename string, src []byte, class, fragmentOk bool) (
 	file *ast.File,
 	sourceAdj func(src []byte, indent int) []byte,
 	indentAdj int,
 	err error,
 ) {
-	file, err = parser.ParseFile(fset, filename, src, parserMode)
+	mode := parserMode
+	if class {
+		mode |= parser.ParseGoPlusClass
+	}
+	file, err = parser.ParseFile(fset, filename, src, mode)
 	return
 }
 

load.go

@@ -51,10 +51,10 @@ func NotFound(err error) bool {
 	return errors.Err(err) == syscall.ENOENT
 }
 
-func loadMod(dir string, gop *env.Gop, conf *Config) (mod *gopmod.Module, err error) {
+func LoadMod(dir string, gop *env.Gop, conf *Config) (mod *gopmod.Module, err error) {
 	mod, err = gopmod.Load(dir, 0)
 	if err != nil && !NotFound(err) {
-		err = errors.NewWith(err, `gopmod.Load(dir, 0`, -2, "gopmod.Load", dir, 0)
+		err = errors.NewWith(err, `gopmod.Load(dir, 0)`, -2, "gopmod.Load", dir, 0)
 		return
 	}
 	if mod != nil {
@@ -101,7 +101,7 @@ func LoadDir(dir string, conf *Config, genTestPkg bool, promptGenGo ...bool) (ou
 	if gop == nil {
 		gop = gopenv.Get()
 	}
-	mod, err := loadMod(dir, gop, conf)
+	mod, err := LoadMod(dir, gop, conf)
 	if err != nil {
 		return
 	}
@@ -177,7 +177,7 @@ func LoadFiles(files []string, conf *Config) (out *gox.Package, err error) {
 	if gop == nil {
 		gop = gopenv.Get()
 	}
-	mod, err := loadMod("", gop, conf)
+	mod, err := LoadMod("", gop, conf)
 	if err != nil {
 		err = errors.NewWith(err, `loadMod("", gop, conf)`, -2, "gop.loadMod", "", gop, conf)
 		return

printer/gop_test.go

@@ -121,7 +121,7 @@ func testFrom(t *testing.T, fpath, sel string, mode int) {
 
 	if (mode & excludeFormatSource) == 0 {
 		t.Run("format.Source "+fpath, func(t *testing.T) {
-			res, err := format.Source(src, fpath)
+			res, err := format.Source(src, false, fpath)
 			if err != nil {
 				t.Fatal("Source failed:", err)
 			}