<commit_before>// Generates methods for enum-like types so they can implement JSON
// marshal/unmarshal and sql valuer/scanner interfaces.
// Borrows ideas and portions of code from github.com/campoy/jsonenums and
// golang.org/x/tools/cmd/stringer

package main

import (
	"bytes"
	"flag"
	"go/ast"
	"go/build"
	"go/constant"
	"go/format"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"html/template"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	var typeNames string
	flag.StringVar(&typeNames, "t", "", "comma-separated list of types")
	flag.Parse()
	if typeNames == "" {
		flag.Usage()
		os.Exit(1)
	}
	typeList := strings.Split(typeNames, ",")
	directory := "."
	if flag.NArg() > 0 {
		directory = flag.Arg(1)
	}
	g := new(Generator)
	g.parsePackageDir(directory)
	for _, typeName := range typeList {
		var analysis = struct {
			Command        string
			PackageName    string
			TypesAndValues map[string][]Value
		}{
			Command:        strings.Join(os.Args, " "),
			PackageName:    g.pkg.name,
			TypesAndValues: make(map[string][]Value),
		}
		g.valuesForType(typeName)
		analysis.TypesAndValues[typeName] = g.values
		var buf bytes.Buffer
		if err := genTemplate.Execute(&buf, analysis); err != nil {
			log.Fatalf("generating code: %v", err)
		}
		src, err := format.Source(buf.Bytes())
		if err != nil {
			log.Printf("internal error: %v", err)
			src = buf.Bytes()
		}
		outputFile := filepath.Join(directory, strings.ToLower(typeName)+"_enum.go")
		if err := ioutil.WriteFile(outputFile, src, 0644); err != nil {
			log.Fatalf("writing output file: %v", err)
		}
	}
}

type Generator struct {
	pkg    *Package
	values []Value
}

type Package struct {
	name  string
	files []*ast.File
	defs  map[*ast.Ident]types.Object
}

func (g *Generator) parsePackageDir(dir string) {
	pkg, err := build.Default.ImportDir(dir, 0)
	if err != nil {
		log.Fatal(err)
	}
	var names []string
	names = append(names, pkg.GoFiles...)
	names = append(names, pkg.CgoFiles...)
	names = prefixDirectory(dir, names)
	g.parsePackage(dir, names)
}

func prefixDirectory(directory string, names []string) []string {
	if directory == "." {
		return names
	}
	ret := make([]string, len(names))
	for i, name := range names {
		ret[i] = filepath.Join(directory, name)
	}
	return ret
}

func (g *Generator) parsePackage(dir string, names []string) {
	var files []*ast.File
	fset := token.NewFileSet()
	for _, name := range names {
		if !strings.HasSuffix(name, ".go") {
			continue
		}
		if dir != "." {
			name = filepath.Join(dir, name)
		}
		f, err := parser.ParseFile(fset, name, nil, 0)
		if err != nil {
			log.Fatal(err)
		}
		files = append(files, f)
	}

	defs := make(map[*ast.Ident]types.Object)
	config := types.Config{Importer: importer.Default()}
	info := &types.Info{Defs: defs}
	if _, err := config.Check(dir, fset, files, info); err != nil {
		log.Printf("type-checking package: %v", err)
	}

	g.pkg = &Package{
		files[0].Name.Name,
		files,
		defs,
	}
}

func (g *Generator) valuesForType(typeName string) {<commit_msg>Light refactor to return errors instead of exiting<commit_after>// Generates methods for enum-like types so they can implement JSON
// marshal/unmarshal and sql valuer/scanner interfaces.
// Borrows ideas and portions of code from github.com/campoy/jsonenums and
// golang.org/x/tools/cmd/stringer

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/constant"
	"go/format"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"html/template"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	var typeNames string
	flag.StringVar(&typeNames, "t", "", "comma-separated list of types")
	flag.Parse()
	if typeNames == "" {
		flag.Usage()
		os.Exit(1)
	}
	typeList := strings.Split(typeNames, ",")
	directory := "."
	if flag.NArg() > 0 {
		directory = flag.Arg(0)
	}
	if err := genEnums(typeList, directory); err != nil {
		log.Fatal(err)
	}
}

func genEnums(typeList []string, dir string) error {
	g := new(Generator)
	if err := g.parsePackageDir(dir); err != nil {
		return err
	}
	for _, typeName := range typeList {
		var analysis = struct {
			Command        string
			PackageName    string
			TypesAndValues map[string][]Value
		}{
			Command:        strings.Join(os.Args, " "),
			PackageName:    g.pkg.name,
			TypesAndValues: make(map[string][]Value),
		}
		g.valuesForType(typeName)
		analysis.TypesAndValues[typeName] = g.values
		var buf bytes.Buffer
		if err := genTemplate.Execute(&buf, analysis); err != nil {
			return err
		}
		src, err := format.Source(buf.Bytes())
		if err != nil {
			log.Printf("WARNING: formatting source: %v", err)
			src = buf.Bytes()
		}
		outputFile := filepath.Join(dir, strings.ToLower(typeName)+"_enum.go")
		if err := ioutil.WriteFile(outputFile, src, 0644); err != nil {
			return fmt.Errorf("writing output file %s: %v", outputFile, err)
		}
	}
	return nil
}

type Generator struct {
	pkg    *Package
	values []Value
}

type Package struct {
	name  string
	files []*ast.File
	defs  map[*ast.Ident]types.Object
}

func (g *Generator) parsePackageDir(dir string) error {
	pkg, err := build.Default.ImportDir(dir, 0)
	if err != nil {
		return fmt.Errorf("importing dir %s: %v", dir, err)
	}
	var names []string
	names = append(names, pkg.GoFiles...)
	names = append(names, pkg.CgoFiles...)
	names = append(names, pkg.SFiles...)
	prefixDirectory(dir, names)
	return g.parsePackage(dir, names)
}

func prefixDirectory(directory string, names []string) {
	if directory == "." {
		return
	}
	for i, name := range names {
		names[i] = filepath.Join(directory, name)
	}
}

func (g *Generator) parsePackage(dir string, names []string) error {
	var files []*ast.File
	fset := token.NewFileSet()
	for _, name := range names {
		if !strings.HasSuffix(name, ".go") {
			continue
		}
		f, err := parser.ParseFile(fset, name, nil, 0)
		if err != nil {
			return fmt.Errorf("parsing file %s: %v", name, err)
		}
		files = append(files, f)
	}

	if len(files) == 0 {
		return fmt.Errorf("%s: no buildable Go files", dir)
	}

	defs := make(map[*ast.Ident]types.Object)
	config := types.Config{Importer: importer.Default()}
	info := &types.Info{Defs: defs}
	if _, err := config.Check(dir, fset, files, info); err != nil {
		return fmt.Errorf("type-checking package: %v", err)
	}

	g.pkg = &Package{
		files[0].Name.Name,
		files,
		defs,
	}

	return nil
}

func (g *Generator) valuesForType(typeName string) {
{'fraction_non_alphanumeric': 0.15714952543955188, 'fraction_numerical': 0.0032674653804263265, 'mean_word_length': 3.6211358734723222, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2151457', 'n_tokens_mistral': 2507, 'n_tokens_neox': 2350, 'n_words': 585}