printer: add generated code warnings

printer/printer.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"runtime"
 	"strings"
 
 	"github.com/mmcloughlin/avo"
@@ -22,13 +23,34 @@ type Config struct {
 }
 
 func NewDefaultConfig() Config {
-	cfg := Config{
+	return Config{
+		Name: "avo",
+		Pkg:  pkg(),
+	}
+}
+
+func NewArgvConfig() Config {
+	return Config{
 		Argv: os.Args,
+		Pkg:  pkg(),
 	}
-	if cwd, err := os.Getwd(); err == nil {
-		cfg.Pkg = filepath.Base(cwd)
+}
+
+// NewGoRunConfig produces a Config for a generator that's expected to be
+// executed via "go run ...".
+func NewGoRunConfig() Config {
+	path := mainfile()
+	if path == "" {
+		return NewDefaultConfig()
+	}
+	argv := []string{"go", "run", filepath.Base(path)}
+	if len(os.Args) > 1 {
+		argv = append(argv, os.Args[1:]...)
+	}
+	return Config{
+		Argv: argv,
+		Pkg:  pkg(),
 	}
-	return cfg
 }
 
 func (c Config) GeneratedBy() string {
@@ -41,3 +63,33 @@ func (c Config) GeneratedBy() string {
 func (c Config) GeneratedWarning() string {
 	return fmt.Sprintf("Code generated by %s. DO NOT EDIT.", c.GeneratedBy())
 }
+
+// mainfile attempts to determine the file path of the main function by
+// inspecting the stack. Returns empty string on failure.
+func mainfile() string {
+	pc := make([]uintptr, 10)
+	n := runtime.Callers(0, pc)
+	if n == 0 {
+		return ""
+	}
+	pc = pc[:n]
+	frames := runtime.CallersFrames(pc)
+	for {
+		frame, more := frames.Next()
+		if frame.Function == "main.main" {
+			return frame.File
+		}
+		if !more {
+			break
+		}
+	}
+	return ""
+}
+
+// pkg guesses the name of the package from the working directory.
+func pkg() string {
+	if cwd, err := os.Getwd(); err == nil {
+		return filepath.Base(cwd)
+	}
+	return ""
+}