go-export

main.go (4745B)

---

 package main
 
 import (
 	"bytes"
 	"fmt"
 	"go/ast"
 	"go/parser"
 	"go/token"
 	"io"
 	"os"
 	"path/filepath"
 	"sort"
 	"strings"
 	"unsafe"
 
 	"github.com/dave/jennifer/jen"
 )
 
 var (
 	// pkgIn is a path to the input package to re-export.
 	pkgIn string
 	// pkgOut is a path to the output package to write the re-exported
 	// symbols.
 	pkgOut string
 	// out is the io object to write to, either a file or stdout.
 	out io.Writer = os.Stdout
 )
 
 func main() {
 	if err := func() (err error) {
 		if len(os.Args) < 1 {
 			return fmt.Errorf("not enough arguments")
 		}
 		args := os.Args[1:]
 		pkgIn = next(&args)
 		pkgOut = next(&args)
 		if pkgOut != "" {
 			if err := os.MkdirAll(pkgOut, 0755); err != nil {
 				return fmt.Errorf("preparing output package: %w", err)
 			}
 			file, err := os.Create(filepath.Join(pkgOut, fmt.Sprintf("%s.go", filepath.Base(pkgIn))))
 			if err != nil {
 				return fmt.Errorf("creating output file: %w", err)
 			}
 			defer file.Close()
 			out = file
 		}
 		return export(pkgIn, out)
 	}(); err != nil {
 		fmt.Printf("error: %v\n", err)
 	}
 }
 
 // next slices off the first element in a string slice.
 func next(args *[]string) string {
 	if args == nil {
 		return ""
 	}
 	if len(*args) == 0 {
 		return ""
 	}
 	defer func() {
 		*args = (*args)[1:]
 	}()
 	return (*args)[0]
 }
 
 // export the in package to an out package.
 func export(in string, out io.Writer) error {
 	in, err := filepath.Abs(in)
 	if err != nil {
 		return fmt.Errorf("resolving path to input package: %w", err)
 	}
 	// Parse the input package.
 	fs := token.NewFileSet()
 	pkgs, err := parser.ParseDir(fs, in, nil, parser.AllErrors)
 	if err != nil {
 		return fmt.Errorf("parsing input package: %w", err)
 	}
 	var (
 		// Assume that the folder name matches the package name - which is a
 		// convention but not a rule.
 		pkgName = filepath.Base(in)
 		// importPath is the fully qualified import path to the target
 		// package we wish to re-export.
 		importPath string
 		seen       = map[string]struct{}{}
 		symbols    []string
 	)
 	importPath, err = buildImportPath(in)
 	if err != nil {
 		return fmt.Errorf("finding module name: %w", err)
 	}
 	target, ok := pkgs[pkgName]
 	if !ok {
 		return fmt.Errorf("package %q not found", pkgName)
 	}
 	// Collect all exported symbols.
 	for _, v := range target.Files {
 		ast.Inspect(v, func(n ast.Node) bool {
 			switch n := n.(type) {
 			case *ast.TypeSpec:
 				if n.Name.IsExported() {
 					_, ok := seen[n.Name.Name]
 					if !ok {
 						symbols = append(symbols, n.Name.Name)
 						seen[n.Name.Name] = struct{}{}
 					}
 				}
 			}
 			return true
 		})
 	}
 	sort.Slice(symbols, func(i, j int) bool {
 		return symbols[i] < symbols[j]
 	})
 	// Generate exported type alias for each exported symbol.
 	f := jen.NewFile(pkgName)
 	f.Type().Defs(func() (c []jen.Code) {
 		for _, s := range symbols {
 			c = append(c, jen.Id(s).Op("=").Qual(importPath, s))
 		}
 		return c
 	}()...)
 	return f.Render(out)
 }
 
 // buildImportPath builds a fully qualified import starting at the package for
 // the path `in`.
 //
 // We does this by walking up directories until we find a go.mod, which we parse
 // for a module name, and join together a fully qualified import path from all
 // the directories we've visited.
 //
 // NOTE(jfm): Assumes that packages names match their directory name - which I
 // don't think is guaranteed.
 func buildImportPath(in string) (string, error) {
 	var (
 		path []string
 		mod  string
 		ok   bool
 	)
 	for {
 		// Record the current fragment to rebuild the path later.
 		path = append(path, filepath.Base(in))
 		// Inspect this directory for a go.mod file.
 		in = filepath.Dir(in)
 		entries, err := os.ReadDir(in)
 		if err != nil {
 			return "", fmt.Errorf("reading directory: %q: %w", in, err)
 		}
 		mod, ok = func() (string, bool) {
 			for _, entry := range entries {
 				if entry.Name() == "go.mod" {
 					return filepath.Join(in, "go.mod"), true
 				}
 			}
 			return "", false
 		}()
 		if ok {
 			break
 		}
 	}
 	// Extract module string from go.mod fie.
 	mod, err := filepath.Abs(mod)
 	if err != nil {
 		return "", fmt.Errorf("resolving go.mod: %w", err)
 	}
 	by, err := os.ReadFile(mod)
 	if err != nil {
 		return "", fmt.Errorf("reading go.mod: %w", err)
 	}
 	parts := bytes.Split(by, []byte{'\n'})
 	if len(parts) == 0 {
 		return "", fmt.Errorf("invalid go.mod")
 	}
 	by = bytes.TrimLeft(bytes.TrimSpace(parts[0]), "module ")
 	// Add module string as final fragment in the path.
 	// Tiny optimization to avoid a string copy, because I can.
 	path = append(path, *(*string)(unsafe.Pointer(&by)))
 	// Reverse the fragments because we generating the slice while unwinding.
 	for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
 		path[i], path[j] = path[j], path[i]
 	}
 	// Join back into a valid import path.
 	return strings.Join(path, "/"), nil
 }