main.go

// Copyright (c) 2020 Cisco and/or its affiliates.
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"text/template"

	"github.com/edwarnicke/exechelper"
)

const (
	importsGoTmpl = `// Code generated by github.com/edwarnicke/imports-gen DO NOT EDIT.
package {{ .Package }}

import (
{{range .Packages}}{{printf "\t_ %q\n" . }}{{end}})
`
)

// TemplateInput - Configuration for imports-gen
type TemplateInput struct {
	// Package - package for generated imports.go
	Package string
	// Packages - packages to import in imports.go
	Packages []string
}

func main() {
	// Load input information from environment
	// go generate passes a number of envs we will pick up this way
	input := &TemplateInput{}
	input.Package = os.Getenv("GOPACKAGE")

	if input.Package == "" {
		log.Fatal("error did not find GOPACKAGE env")
	}

	// Setup default exechelper.Options
	options := []*exechelper.Option{
		exechelper.WithEnvirons(os.Environ()...),
		exechelper.WithStderr(os.Stderr),
	}

	// Extract the GOOS
	goos, ok := os.LookupEnv("GOOS")
	if !ok {
		cmdStr := "go env GOOS"
		goosBytes, GoosErr := exechelper.Output(cmdStr, options...)
		if GoosErr != nil {
			log.Fatalf("error extracting GOOS: %+v", GoosErr)
		}
		goos = strings.TrimSpace(string(goosBytes))
	}

	// Compute the filename, and remove any old one *before* we build up the list of packages we need
	filename := fmt.Sprintf("imports_%s.go", goos)
	if err := os.Remove(filename); err != nil && !os.IsNotExist(err) {
		log.Fatalf("unable to remove %s because %+v", filename, err)
	}

	// Get the module Path and Dir (example: Path: github.com/foo/bar Dir: /home/bob/git/bar)
	modBytes, err := exechelper.Output(`go list -m -f '{{printf "%s %s" .Path .Dir}}'`, options...)
	if err != nil {
		log.Fatalf("error exctracting Module Directory: %+v", err)
	}
	if len(modBytes) == 0 {
		log.Fatal("error extracting Module Directory, no module directory found.  import-gen can only be used with go modules")
	}
	mods := strings.Split(strings.TrimSpace(string(modBytes)), " ")
	modPath := mods[0]
	modDir := mods[1]

	// Get the imported packages for the module
	options = append(options, exechelper.WithDir(modDir))
	packagesBytes, err := exechelper.Output("go list -f '{{range .Imports}}{{printf \"%s\\n\" .}}{{end}}{{range .XTestImports}}{{printf \"%s\\n\" .}}{{end}}{{range .TestImports}}{{printf \"%s\\n\" .}}{{end}}' ./...", options...)
	if err != nil {
		log.Fatalf("error exctracting packages: %+v", err)
	}
	if len(packagesBytes) == 0 {
		log.Fatal("error extracting packages, no packages found.  import-gen can only be used if a module imports packages")
	}

	rawPackages := strings.Split(strings.TrimSpace(string(packagesBytes)), "\n")
	packagesSeen := make(map[string]bool)
	for _, pkg := range rawPackages {
		if !(strings.HasPrefix(pkg, modPath) || packagesSeen[pkg] || (pkg == "C")) {
			input.Packages = append(input.Packages, pkg)
			packagesSeen[pkg] = true
		}
	}
	sort.Strings(input.Packages)

	// Create the template
	tmpl := template.Must(template.New(fmt.Sprintf("imports_%s.go", goos)).Parse(importsGoTmpl))

	// Create the imports.go file
	f, err := os.Create(filename)
	if err != nil {
		log.Fatalf("error creating file: %q: %+v", filename, err)
	}
	defer func() { _ = f.Close() }()
	if err := tmpl.Execute(f, input); err != nil {
		fmt.Fprintf(os.Stderr, "error processing template: %+v", err)
	}
}