<commit_before>// Copyright 2012 <NAME>. All rights reserved.
// Use of this source code is governed by a 3-clause BSD
// license that can be found in the LICENSE file.

package main

// TODO:
// - Default arguments. Possibly don't bother...

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/samuel/go-thrift/parser"
)

var (
	flagGoBinarystring = flag.Bool("go.binarystring", false, "Always use string for binary instead of []byte")
	flagGoJsonEnumnum  = flag.Bool("go.json.enumnum", false, "For JSON marshal enums by number instead of name")
	flagGoPointers     = flag.Bool("go.pointers", false, "Make all fields pointers")
)

var (
	goNamespaceOrder = []string{"go", "perl", "py", "cpp"}
)

type ErrUnknownType string

func (e ErrUnknownType) Error() string {
	return fmt.Sprintf("Unknown type %s", string(e))
}

type ErrMissingInclude string

func (e ErrMissingInclude) Error() string {
	return fmt.Sprintf("Missing include %s", string(e))
}

type GoGenerator struct {
	thrift *parser.Thrift
	pkg    string

	ThriftFiles map[string]*parser.Thrift
	Packages    map[string]string
}

func (g *GoGenerator) error(err error) {
	panic(err)
}

func (g *GoGenerator) write(w io.Writer, f string, a ...interface{}) error {
	if _, err := io.WriteString(w, fmt.Sprintf(f, a...)); err != nil {
		g.error(err)
		return err
	}
	return nil
}

func (g *GoGenerator) formatType(pkg string, thrift *parser.Thrift, typ *parser.Type, optional bool) string {
	if strings.Contains(typ.Name, ".") {
		parts := strings.SplitN(typ.Name, ".", 2)
		thriftFilename := thrift.Includes[parts[0]]
		if thriftFilename == "" {
			g.error(ErrMissingInclude(parts[0]))
		}
		thrift = g.ThriftFiles[thriftFilename]
		if thrift == nil {
			g.error(ErrMissingInclude(thriftFilename))
		}
		pkg = g.Packages[thriftFilename]
		typ = &parser.Type{
			Name:      parts[1],
			KeyType:   typ.KeyType,
			ValueType: typ.ValueType,
		}
	}

	ptr := ""
	if *flagGoPointers || optional {
		ptr = "*"
	}
	switch typ.Name {
	case "byte", "bool", "string":
		return ptr + typ.Name
	case "binary":
		if *flagGoBinarystring {
			return ptr + "string"
		}
		return "[]byte"
	case "i16":
		return ptr + "int16"
	case "i32":
		return ptr + "int32"
	case "i64":
		return ptr + "int64"
	case "double":
		return ptr + "float64"
	case "set":
		valueType := g.formatType(pkg, thrift, typ.ValueType, false)
		if valueType == "[]byte" {
			valueType = "string"
		}
		return "map[" + valueType + "]struct{}"
	case "list":
		return "[]" + g.formatType(pkg, thrift, typ.ValueType, false)
	case "map":
		keyType := g.formatType(pkg, thrift, typ.KeyType, false)
		if keyType == "[]byte" {
			// TODO: Log, warn, do something!
			// println("key type of []byte not supported for maps")
			keyType = "string"
		}
		return "map[" + keyType + "]" + g.formatType(pkg, thrift, typ.ValueType, false)
	}

	if t := thrift.Typedefs[typ.Name]; t != nil {
		return g.formatType(pkg, thrift, t, optional)
	}
	if e := thrift.Enums[typ.Name]; e != nil {
		name := e.Name
		if pkg != g.pkg {
			name = pkg + "." + name
		}
		return ptr + name
	}
	if s := thrift.Structs[typ.Name]; s != nil {
		name := s.Name
		if pkg != g.pkg {
			name = pkg + "." + name
		}
		return "*" + name
	}
	if e := thrift.Exceptions[typ.Name]; e != nil {
		name := e.Name
		if pkg != g.pkg {
			name = pkg + "." + name
		}
		return "*" + name
	}

	g.error(ErrUnknownType(typ.Name))
	return ""
}

func (g *GoGenerator) formatField(field *parser.Field) string {
	tags := ""
	if !field.Optional {
		tags = ",required"
	}
	return fmt.Sprintf(
		"%s %s `thrift:\"%d%s\" json:\"%s\"`",
		camelCase(field.Name), g.formatType(g.pkg, g.thrift, field.Type, field.Optional), field.Id, tags, field.Name)
}

func (g *GoGenerator) formatArguments(arguments []*parser.Field) string {
	args := make([]string, len(arguments))
	for i, arg := range arguments {
		args[i] = fmt.Sprintf("%s %s", camelCase(arg.Name), g.formatType(g.pkg, g.thrift, arg.Type, arg.Optional))
	}
	return strings.Join(args, ", ")
}

func (g *GoGenerator) formatReturnType(typ *parser.Type) string {
	if typ == nil || typ.Name == "void" {
		return "error"
	}
	return fmt.Sprintf("(%s, error)", g.formatType(g.pkg, g.thrift, typ, false))
}

func (g *GoGenerator) writeConstant(out io.Writer, c *parser.Constant) error {
	return g.write(out, "\nconst %s = %+v\n", camelCase(c.Name), c.Value)
}

func (g *GoGenerator) writeEnum(out io.Writer, enum *parser.Enum) error {
	enumName := camelCase(enum.Name)<commit_msg>Add omitempty to json tag for optional fields<commit_after>// Copyright 2012 <NAME>. All rights reserved.
// Use of this source code is governed by a 3-clause BSD
// license that can be found in the LICENSE file.

package main

// TODO:
// - Default arguments. Possibly don't bother...

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/samuel/go-thrift/parser"
)

var (
	flagGoBinarystring = flag.Bool("go.binarystring", false, "Always use string for binary instead of []byte")
	flagGoJsonEnumnum  = flag.Bool("go.json.enumnum", false, "For JSON marshal enums by number instead of name")
	flagGoPointers     = flag.Bool("go.pointers", false, "Make all fields pointers")
)

var (
	goNamespaceOrder = []string{"go", "perl", "py", "cpp"}
)

type ErrUnknownType string

func (e ErrUnknownType) Error() string {
	return fmt.Sprintf("Unknown type %s", string(e))
}

type ErrMissingInclude string

func (e ErrMissingInclude) Error() string {
	return fmt.Sprintf("Missing include %s", string(e))
}

type GoGenerator struct {
	thrift *parser.Thrift
	pkg    string

	ThriftFiles map[string]*parser.Thrift
	Packages    map[string]string
}

func (g *GoGenerator) error(err error) {
	panic(err)
}

func (g *GoGenerator) write(w io.Writer, f string, a ...interface{}) error {
	if _, err := io.WriteString(w, fmt.Sprintf(f, a...)); err != nil {
		g.error(err)
		return err
	}
	return nil
}

func (g *GoGenerator) formatType(pkg string, thrift *parser.Thrift, typ *parser.Type, optional bool) string {
	if strings.Contains(typ.Name, ".") {
		parts := strings.SplitN(typ.Name, ".", 2)
		thriftFilename := thrift.Includes[parts[0]]
		if thriftFilename == "" {
			g.error(ErrMissingInclude(parts[0]))
		}
		thrift = g.ThriftFiles[thriftFilename]
		if thrift == nil {
			g.error(ErrMissingInclude(thriftFilename))
		}
		pkg = g.Packages[thriftFilename]
		typ = &parser.Type{
			Name:      parts[1],
			KeyType:   typ.KeyType,
			ValueType: typ.ValueType,
		}
	}

	ptr := ""
	if *flagGoPointers || optional {
		ptr = "*"
	}
	switch typ.Name {
	case "byte", "bool", "string":
		return ptr + typ.Name
	case "binary":
		if *flagGoBinarystring {
			return ptr + "string"
		}
		return "[]byte"
	case "i16":
		return ptr + "int16"
	case "i32":
		return ptr + "int32"
	case "i64":
		return ptr + "int64"
	case "double":
		return ptr + "float64"
	case "set":
		valueType := g.formatType(pkg, thrift, typ.ValueType, false)
		if valueType == "[]byte" {
			valueType = "string"
		}
		return "map[" + valueType + "]struct{}"
	case "list":
		return "[]" + g.formatType(pkg, thrift, typ.ValueType, false)
	case "map":
		keyType := g.formatType(pkg, thrift, typ.KeyType, false)
		if keyType == "[]byte" {
			// TODO: Log, warn, do something!
			// println("key type of []byte not supported for maps")
			keyType = "string"
		}
		return "map[" + keyType + "]" + g.formatType(pkg, thrift, typ.ValueType, false)
	}

	if t := thrift.Typedefs[typ.Name]; t != nil {
		return g.formatType(pkg, thrift, t, optional)
	}
	if e := thrift.Enums[typ.Name]; e != nil {
		name := e.Name
		if pkg != g.pkg {
			name = pkg + "." + name
		}
		return ptr + name
	}
	if s := thrift.Structs[typ.Name]; s != nil {
		name := s.Name
		if pkg != g.pkg {
			name = pkg + "." + name
		}
		return "*" + name
	}
	if e := thrift.Exceptions[typ.Name]; e != nil {
		name := e.Name
		if pkg != g.pkg {
			name = pkg + "." + name
		}
		return "*" + name
	}

	g.error(ErrUnknownType(typ.Name))
	return ""
}

func (g *GoGenerator) formatField(field *parser.Field) string {
	tags := ""
	jsonTags := ""
	if !field.Optional {
		tags = ",required"
	} else {
		jsonTags = ",omitempty"
	}
	return fmt.Sprintf(
		"%s %s `thrift:\"%d%s\" json:\"%s%s\"`",
		camelCase(field.Name), g.formatType(g.pkg, g.thrift, field.Type, field.Optional), field.Id, tags, field.Name, jsonTags)
}

func (g *GoGenerator) formatArguments(arguments []*parser.Field) string {
	args := make([]string, len(arguments))
	for i, arg := range arguments {
		args[i] = fmt.Sprintf("%s %s", camelCase(arg.Name), g.formatType(g.pkg, g.thrift, arg.Type, arg.Optional))
	}
	return strings.Join(args, ", ")
}

func (g *GoGenerator) formatReturnType(typ *parser.Type) string {
	if typ == nil || typ.Name == "void" {
		return "error"
	}
	return fmt.Sprintf("(%s, error)", g.formatType(g.pkg, g.thrift, typ, false))
}

func (g *GoGenerator) writeConstant(out io.Writer, c *parser.Constant) error {
	return g.write(out, "\nconst %s = %+v\n", camelCase(c.Name), c.Value)
}

func (g *GoGenerator) writeEnum(out io.Writer, enum *parser.Enum) error {
	enumName := camelCase(enum.Name)
{'fraction_non_alphanumeric': 0.1678688524590164, 'fraction_numerical': 0.005027322404371585, 'mean_word_length': 3.659368635437882, 'pattern_counts': {'":': 18, '<': 5, '<?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': '6090993', 'n_tokens_mistral': 3562, 'n_tokens_neox': 3359, 'n_words': 883}