kustomize/kyaml/yaml/internal/k8sgen/generator/main.go

// Copyright 2020 The Kubernetes Authors.
// SPDX-License-Identifier: Apache-2.0

// A code copier launched via `go generate`.
// See k8segen/doc.go for further discussion.
package main

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const (
	// Splitting this so it doesn't show up in grep.
	apiMachineryModule  = "k8s.io" + "/" + "apimachinery"
	apiMachineryVersion = "v0.17.0"
	sigsK8sIo           = "sigs.k8s.io"
)

var (
	filesToCopy = map[string][]string{
		filepath.Join("pkg", "labels"): {
			"labels.go",
			"selector.go",
			"zz_generated.deepcopy.go",
		},
		filepath.Join("pkg", "selection"): {
			"operator.go",
		},
		filepath.Join("pkg", "util", "sets"): {
			"empty.go",
			"string.go",
		},
		filepath.Join("pkg", "util", "errors"): {
			"errors.go",
		},
		filepath.Join("pkg", "util", "validation"): {
			"validation.go",
		},
		filepath.Join("pkg", "util", "validation", "field"): {
			"errors.go",
			"path.go",
		},
	}
)

func main() {
	c := newCopier()
	// c.print()
	runNoOutputCommand("go", "get", c.apiMachineryModSpec())
	for dir, files := range filesToCopy {
		for _, n := range files {
			if err := c.copyFile(dir, n); err != nil {
				log.Fatal(err)
			}
		}
	}
	runNoOutputCommand(
		"go", "mod", "edit", "-droprequire="+apiMachineryModule)
	runNoOutputCommand("go", "mod", "tidy")
	runGetOutputCommand("go", "fmt", "./...")
}

type copier struct {
	goModCache string
	topPackage string
	srcDir     string
	pgmName    string
}

func (c copier) apiMachineryModSpec() string {
	return apiMachineryModule + "@" + apiMachineryVersion
}

func (c copier) replacementPath() string {
	return filepath.Join(c.topPackage, c.subPath())
}

func (c copier) subPath() string {
	return filepath.Join("internal", c.pgmName)
}

func (c copier) print() {
	fmt.Printf("    apiMachineryModule: %s\n", apiMachineryModule)
	fmt.Printf("     replacementPath: %s\n", c.replacementPath())
	fmt.Printf("          goModCache: %s\n", c.goModCache)
	fmt.Printf("          topPackage: %s\n", c.topPackage)
	fmt.Printf("             subPath: %s\n", c.subPath())
	fmt.Printf("              srcDir: %s\n", c.srcDir)
	fmt.Printf("  apiMachineryModSpec: %s\n", c.apiMachineryModSpec())
	fmt.Printf("             pgmName: %s\n", c.pgmName)
	fmt.Printf("                 pwd: %s\n", os.Getenv("PWD"))
}

func newCopier() copier {
	tmp := copier{
		pgmName:    os.Getenv("GOPACKAGE"),
		goModCache: runGetOutputCommand("go", "env", "GOMODCACHE"),
	}
	goMod := runGetOutputCommand("go", "env", "GOMOD")
	topPackage := filepath.Join(goMod[:len(goMod)-len("go.mod")-1], "yaml")
	k := strings.Index(topPackage, sigsK8sIo)
	if k < 1 {
		log.Fatalf("cannot find %s in %s", sigsK8sIo, topPackage)
	}
	tmp.srcDir = topPackage[:k-1]
	tmp.topPackage = topPackage[k:]
	return tmp
}

func (c copier) copyFile(dir, name string) error {
	inFile, err := os.Open(
		filepath.Join(c.goModCache, c.apiMachineryModSpec(), dir, name))
	if err != nil {
		return err
	}
	defer inFile.Close()
	scanner := bufio.NewScanner(inFile)

	w, err := newWriter(dir, name)
	if err != nil {
		return err
	}
	defer w.close()

	w.write(
		fmt.Sprintf(
			// This particular phrasing is required.
			"// Code generated by %s/generator from %s; DO NOT EDIT.",
			c.pgmName, c.apiMachineryModSpec()))
	w.write(
		fmt.Sprintf(
			"// Copied from %s\n",
			filepath.Join(c.apiMachineryModSpec(), dir, name)))

	for scanner.Scan() {
		l := scanner.Text()
		// Disallow recursive generation.
		if strings.HasPrefix(l, "//go:generate") {
			continue
		}
		// Don't want it to appear double generated.
		if strings.HasPrefix(l, "// Code generated") {
			continue
		}
		// Fix self-imports.
		l = strings.Replace(l, apiMachineryModule, c.replacementPath(), 1)
		// Replace klog with generic log (eschewing k8s.io entirely).
		l = strings.Replace(l, "\"k8s.io/klog\"", "\"log\"", 1)
		l = strings.Replace(l, "klog.V(10).Infof(", "log.Printf(", 1)
		w.write(l)
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	w.write("")
	return nil
}

type writer struct {
	root string
	f    *os.File
}

func newWriter(toDir, name string) (*writer, error) {
	if err := os.MkdirAll(toDir, 0755); err != nil {
		log.Printf("unable to create directory: %s", toDir)
		return nil, err
	}
	n := filepath.Join(toDir, name)
	f, err := os.Create(n)
	if err != nil {
		return nil, fmt.Errorf("unable to create `%s`; %v", n, err)
	}
	return &writer{root: toDir, f: f}, nil
}

func (w *writer) close() {
	w.f.Close()
}

func (w *writer) write(line string) {
	if _, err := w.f.WriteString(line + "\n"); err != nil {
		log.Printf("Trouble writing: %s", line)
		log.Fatalf("Error: %s", err)
	}
}

func runNoOutputCommand(n string, args ...string) {
	o := runGetOutputCommand(n, args...)
	if len(o) > 0 {
		log.Fatalf("unexpected output: %q", o)
	}
}

func runGetOutputCommand(n string, args ...string) string {
	cmd := exec.Command(n, args...)
	var outBuf bytes.Buffer
	cmd.Stdout = &outBuf
	var errBuf bytes.Buffer
	cmd.Stderr = &errBuf
	if err := cmd.Run(); err != nil {
		fmt.Printf("err: %q\n", errBuf.String())
		log.Fatal(err)
	}
	return strings.TrimSpace(outBuf.String())
}