go/types: allow filename arguments to generator for easier manual use

Change-Id: I96b25fae1dcbfae55f57c07de9529c6663ad28d4
Reviewed-on: https://go-review.googlesource.com/c/go/+/461602
Reviewed-by: Robert Griesemer <gri@google.com>
Auto-Submit: Robert Griesemer <gri@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
Run-TryBot: Robert Griesemer <gri@google.com>

diff --git a/src/go/types/generator.go b/src/go/types/generator.go
index 4f7d6661ef1bbd6d29076962bdd83fb6d776ae39..2972f5e1bfaee953578bb475aab841debaa00898 100644 (file)
--- a/src/go/types/generator.go
+++ b/src/go/types/generator.go
@@ -11,6 +11,8 @@ package main
 
 import (
        "bytes"
+       "flag"
+       "fmt"
        "go/ast"
        "go/format"
        "go/parser"
@@ -30,34 +32,50 @@ const (
 var fset = token.NewFileSet()
 
 func main() {
-       for filename, action := range filemap {
-               // parse src
-               srcFilename := filepath.FromSlash(runtime.GOROOT() + "/src/" + srcDir + "/" + filename)
-               file, err := parser.ParseFile(fset, srcFilename, nil, parser.ParseComments)
-               if err != nil {
-                       log.Fatal(err)
+       flag.Parse()
+
+       // process provided filenames, if any
+       if flag.NArg() > 0 {
+               for _, filename := range flag.Args() {
+                       fmt.Println("generating", filename)
+                       generate(filename, filemap[filename])
                }
+               return
+       }
 
-               // fix package name
-               file.Name.Name = strings.ReplaceAll(file.Name.Name, "types2", "types")
+       // otherwise process per filemap below
+       for filename, action := range filemap {
+               generate(filename, action)
+       }
+}
 
-               // rewrite AST as needed
-               if action != nil {
-                       action(file)
-               }
+func generate(filename string, action action) {
+       // parse src
+       srcFilename := filepath.FromSlash(runtime.GOROOT() + "/src/" + srcDir + "/" + filename)
+       file, err := parser.ParseFile(fset, srcFilename, nil, parser.ParseComments)
+       if err != nil {
+               log.Fatal(err)
+       }
 
-               // format AST
-               var buf bytes.Buffer
-               buf.WriteString("// Code generated by \"go run generator.go\"; DO NOT EDIT.\n\n")
-               if err := format.Node(&buf, fset, file); err != nil {
-                       log.Fatal(err)
-               }
+       // fix package name
+       file.Name.Name = strings.ReplaceAll(file.Name.Name, "types2", "types")
 
-               // write dst
-               dstFilename := filepath.FromSlash(runtime.GOROOT() + "/src/" + dstDir + "/" + filename)
-               if err := os.WriteFile(dstFilename, buf.Bytes(), 0o644); err != nil {
-                       log.Fatal(err)
-               }
+       // rewrite AST as needed
+       if action != nil {
+               action(file)
+       }
+
+       // format AST
+       var buf bytes.Buffer
+       buf.WriteString("// Code generated by \"go run generator.go\"; DO NOT EDIT.\n\n")
+       if err := format.Node(&buf, fset, file); err != nil {
+               log.Fatal(err)
+       }
+
+       // write dst
+       dstFilename := filepath.FromSlash(runtime.GOROOT() + "/src/" + dstDir + "/" + filename)
+       if err := os.WriteFile(dstFilename, buf.Bytes(), 0o644); err != nil {
+               log.Fatal(err)
        }
 }