genai: switch gen-examples flow to use go:generate (#168)

Follow up on #164 --------- Co-authored-by: Jonathan Amsterdam <[email protected]>
google · Jul 11, 2024 · 3d279eb · 3d279eb
1 parent e91ac8d
commit 3d279eb
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 17 deletions.
diff --git a/genai/example_test.go b/genai/example_test.go
@@ -1,5 +1,4 @@
-// DO NOT EDIT THIS FILE -- it is automatically generated.
-// See internal/cmd/gen-examples for details.
+// This file was generated from internal/snippets/example_test.go. DO NOT EDIT.
 
 // Copyright 2023 Google LLC
 //

diff --git a/genai/internal/cmd/gen-examples/gen-examples.go b/genai/internal/cmd/gen-examples/gen-examples.go
@@ -16,14 +16,12 @@
 // and copies them to "official" examples in genai/example_test.go, while
 // removing snippet comments (between [START...] and [END...]) that are used
 // for website documentation purposes.
-//
-// Run from the root directory:
-//
-//   go run ./genai/internal/cmd/gen-examples/ < genai/internal/snippets/example_test.go > genai/example_test.go
+// It's invoked with a go:generate directive in the source file.
 
 package main
 
 import (
+	"flag"
 	"fmt"
 	"go/ast"
 	"go/format"
@@ -35,8 +33,22 @@ import (
 )
 
 func main() {
+	inPath := flag.String("in", "", "input file path")
+	outPath := flag.String("out", "", "output file path")
+	flag.Parse()
+
+	if len(*inPath) == 0 || len(*outPath) == 0 {
+		log.Fatalf("got empty -in (%v) or -out (%v)", *inPath, *outPath)
+	}
+
+	inFile, err := os.Open(*inPath)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer inFile.Close()
+
 	fset := token.NewFileSet()
-	file, err := parser.ParseFile(fset, "src.go", os.Stdin, parser.ParseComments)
+	file, err := parser.ParseFile(fset, *inPath, inFile, parser.ParseComments)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -45,13 +57,18 @@ func main() {
 		sanitizeCommentGroup(cgroup)
 	}
 
-	fmt.Println(strings.TrimLeft(preamble, "\r\n"))
-	format.Node(os.Stdout, fset, file)
+	outFile, err := os.Create(*outPath)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer outFile.Close()
+
+	fmt.Fprintln(outFile, strings.TrimLeft(preamble, "\r\n"))
+	format.Node(outFile, fset, file)
 }
 
 const preamble = `
-// DO NOT EDIT THIS FILE -- it is automatically generated.
-// See internal/cmd/gen-examples for details.
+// This file was generated from internal/snippets/example_test.go. DO NOT EDIT.
 `
 
 func printCommentGroup(cg *ast.CommentGroup) {
@@ -62,17 +79,20 @@ func printCommentGroup(cg *ast.CommentGroup) {
 }
 
 // sanitizeCommentGroup removes comment blocks between [START... and [END...
-// (including these lines) - it modifies cg.
+// (including these lines), and also any go:generate directives - it modifies cg.
 func sanitizeCommentGroup(cg *ast.CommentGroup) {
 	var nl []*ast.Comment
-	exclude := false
+	excludeBlock := false
 	for _, commentLine := range cg.List {
 		if strings.Contains(commentLine.Text, "[START") {
-			exclude = true
+			excludeBlock = true
 		} else if strings.Contains(commentLine.Text, "[END") {
-			exclude = false
-		} else if !exclude {
-			nl = append(nl, commentLine)
+			excludeBlock = false
+		} else if !excludeBlock {
+
+			if !strings.Contains(commentLine.Text, "go:generate") {
+				nl = append(nl, commentLine)
+			}
 		}
 	}
 	cg.List = nl

diff --git a/genai/internal/snippets/example_test.go b/genai/internal/snippets/example_test.go
@@ -14,6 +14,8 @@
 
 package genai_test
 
+//go:generate go run ../cmd/gen-examples/gen-examples.go -in $GOFILE -out ../../example_test.go
+
 import (
 	"context"
 	"encoding/json"