cmd/gotidy/main.go

// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"fuchsia.googlesource.com/infra/infra/gotidy"
)

// TidyStep transforms source code into other source code.
type TidyStep interface {
	Tidy([]byte) ([]byte, error)
}

// TidyFunc is an adapter for functions with the same signature as TidyStep.Tidy.
type TidyFunc func([]byte) ([]byte, error)

// Tidy implements TidyStep.
func (fn TidyFunc) Tidy(input []byte) ([]byte, error) {
	return fn(input)
}

var (
	// Whether to overwrite the input file contents with tidied output.
	overwrite bool

	// Whether to display usage and exit.
	help bool

	// The list of tidy steps to run.
	tidySteps = []TidyStep{
		&gotidy.AddCopyrightHeader{Year: time.Now().Year()},
		TidyFunc(format.Source), // gofmt. Always do this last.
	}
)

func usage() {
	fmt.Fprint(os.Stderr, "gotidy [flags] [path...]")
	flag.PrintDefaults()
}

func init() {
	flag.BoolVar(&help, "h", false, "Display usage and exit")
	flag.BoolVar(&overwrite, "w", false, "Whether to overwrite files with tidied output")
	flag.Usage = usage
}

func main() {
	flag.Parse()

	if help || flag.NArg() == 0 {
		flag.Usage()
		return
	}

	os.Exit(execute(flag.Args()))
}

func execute(paths []string) int {
	if len(paths) == 0 {
		paths = []string{"."}
	}

	var files []string
	for _, path := range paths {
		files = append(files, listGoFilesRecursive(path)...)
	}

	return exitCode(tidyFiles(files))
}

func listGoFilesRecursive(root string) []string {
	var paths []string
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() && path != root {
			paths = append(paths, listGoFilesRecursive(path)...)
		}

		if strings.HasSuffix(path, ".go") {
			paths = append(paths, path)
		}
		return nil
	})

	if err != nil {
		return []string{}
	}

	return paths
}

func tidyFiles(files []string) <-chan error {
	errs := make(chan error)

	var wg sync.WaitGroup
	for _, file := range files {
		wg.Add(1)
		log.Println("tidying", file)
		go tidyFile(file, errs, &wg)
	}

	go func() {
		wg.Wait()
		close(errs)
	}()

	return errs
}

func tidyFile(file string, errs chan<- error, wg *sync.WaitGroup) {
	input, err := ioutil.ReadFile(file)
	if err != nil {
		errs <- fmt.Errorf("failed to read %s: %v", file, err)
	}

	tidied, err := tidy(input)
	if err != nil {
		errs <- fmt.Errorf("failed to tidy %s: %v", file, err)
	}

	reader := bytes.NewReader(tidied)
	dest := os.Stdout
	if overwrite {
		dest, err = os.Create(file)
		if err != nil {
			errs <- fmt.Errorf("failed to open %s: %v", file, err)
		}
	}

	if _, err := io.Copy(dest, reader); err != nil {
		errs <- fmt.Errorf("failed to write output to %s: %v", file, err)
	}

	wg.Done()
}

func tidy(input []byte) ([]byte, error) {
	source := input

	for _, tidyStep := range tidySteps {
		output, err := tidyStep.Tidy(source)
		if err != nil {
			return nil, err
		}
		source = output
	}

	return source, nil
}

func exitCode(in <-chan error) int {
	var code int
	for err := range in {
		log.Println(err)
		code = 1
	}
	return code
}