internal/cmd/weave: add -o flag

Add a -o flag to the weave command, so that it can be more easily used
with go:generate.

Change-Id: I620c99528ab8156dd89700b1576c445c6ac7c55a
Reviewed-on: https://go-review.googlesource.com/c/example/+/702377
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Auto-Submit: Robert Findley <rfindley@google.com>

internal/cmd/weave/weave.go

diff --git a/internal/cmd/weave/weave.go b/internal/cmd/weave/weave.go
index 880b415..d20e03c 100644
--- a/internal/cmd/weave/weave.go
+++ b/internal/cmd/weave/weave.go
@@ -37,19 +37,21 @@
 	"strings"
 )
 
+var output = flag.String("o", "", "output file (empty means stdout)")
+
 func main() {
+	flag.Usage = func() {
+		fmt.Fprintf(flag.CommandLine.Output(), "usage: weave [flags] <input.md>\n\nflags:\n")
+		flag.PrintDefaults()
+	}
 	flag.Parse()
-	log.SetFlags(0)
-	log.SetPrefix("weave: ")
 	if flag.NArg() != 1 {
-		log.Fatal("usage: weave input.md\n")
+		flag.Usage()
+		os.Exit(2)
 	}
 
-	f, err := os.Open(flag.Arg(0))
-	if err != nil {
-		log.Fatal(err)
-	}
-	defer f.Close()
+	log.SetFlags(0)
+	log.SetPrefix("weave: ")
 
 	wd, err := os.Getwd()
 	if err != nil {
@@ -57,13 +59,38 @@
 	}
 	curDir := filepath.Base(wd)
 
-	fmt.Println("<!-- Autogenerated by weave; DO NOT EDIT -->")
+	in, err := os.Open(flag.Arg(0))
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer in.Close()
+
+	out := os.Stdout
+	if *output != "" {
+		out, err = os.Create(*output)
+		if err != nil {
+			log.Fatal(err)
+		}
+		defer func() {
+			if err := out.Close(); err != nil {
+				log.Fatal(err)
+			}
+		}()
+	}
+
+	printf := func(format string, args ...any) {
+		if _, err := fmt.Fprintf(out, format, args...); err != nil {
+			log.Fatalf("writing failed: %v", err)
+		}
+	}
+
+	printf("<!-- Autogenerated by weave; DO NOT EDIT -->")
 
 	// Pass 1: extract table of contents.
 	var toc []string
-	in := bufio.NewScanner(f)
-	for in.Scan() {
-		line := in.Text()
+	scanner := bufio.NewScanner(in)
+	for scanner.Scan() {
+		line := scanner.Text()
 		if line == "" || (line[0] != '#' && line[0] != '%') {
 			continue
 		}
@@ -84,21 +111,21 @@
 			toc = append(toc, line)
 		}
 	}
-	if in.Err() != nil {
-		log.Fatal(in.Err())
+	if scanner.Err() != nil {
+		log.Fatal(scanner.Err())
 	}
 
 	// Pass 2.
-	if _, err := f.Seek(0, io.SeekStart); err != nil {
+	if _, err := in.Seek(0, io.SeekStart); err != nil {
 		log.Fatalf("can't rewind input: %v", err)
 	}
-	in = bufio.NewScanner(f)
-	for in.Scan() {
-		line := in.Text()
+	scanner = bufio.NewScanner(in)
+	for scanner.Scan() {
+		line := scanner.Text()
 		switch {
 		case strings.HasPrefix(line, "%toc"): // ToC
 			for _, h := range toc {
-				fmt.Println(h)
+				printf("%s\n", h)
 			}
 		case strings.HasPrefix(line, "%include"):
 			words := strings.Fields(line)
@@ -125,7 +152,7 @@
 			filename := words[1]
 
 			if caption {
-				fmt.Printf("	// go get golang.org/x/example/%s/%s\n\n",
+				printf("	// go get golang.org/x/example/%s/%s\n\n",
 					curDir, filepath.Dir(filename))
 			}
 
@@ -133,15 +160,15 @@
 			if err != nil {
 				log.Fatal(err)
 			}
-			fmt.Println("```go")
-			fmt.Println(cleanListing(s)) // TODO(adonovan): escape /^```/ in s
-			fmt.Println("```")
+			printf("```go\n")
+			printf("%s\n", cleanListing(s)) // TODO(adonovan): escape /^```/ in s
+			printf("```\n")
 		default:
-			fmt.Println(line)
+			printf("%s\n", line)
 		}
 	}
-	if in.Err() != nil {
-		log.Fatal(in.Err())
+	if scanner.Err() != nil {
+		log.Fatal(scanner.Err())
 	}
 }