csvgen.go

package main

import (
	"fmt"
	. "github.com/dave/jennifer/jen"
	"github.com/karantin2020/cli"
	"github.com/karantin2020/csvgen/parser"
	"os"
	"path/filepath"
	"strings"
)

var (
	// Config vars
	pkg        string
	subpkg     string
	out        string
	fname      string
	parseEmpty bool
	marshal    bool
	unmarshal  bool

	fInfo os.FileInfo

	// Package vars
	pkgCnt string
	f      *File
	p      = parser.Parser{AllStructs: true}
)

const fieldPrefix = "um"

func main() {
	flags := cli.New("This app generates csv Marshall and Unmarshal functions", "0.1.2")
	flags.StringVarP(&pkg, "pkg", "p", "", "output package")
	flags.StringVarP(&subpkg, "subpkg", "s", "", "output subpkg name")
	flags.StringVarP(&fname, "fname", "f", "", "input file")
	flags.StringVarP(&out, "out", "o", "", "output file")
	flags.BoolVarP(&parseEmpty, "parseempty", "e", false, "parse empty fields with values: '', 0, 0.0, false. Default: false")
	flags.BoolVarP(&marshal, "marshal", "m", true, "generate MarshalCSV or not")
	flags.BoolVarP(&unmarshal, "unmarshal", "u", true, "generate UnmarshalCSV or not")

	flags.Parse()

	pkgCnt = "main"
	if pkg != "" {
		pkgCnt = pkg
	}
	if subpkg != "" {
		pkgCnt = subpkg
		pkg = ""
	}

	if fname == "" {
		fmt.Println("didn't pass source file to parse")
		os.Exit(1)
	}
	var err error
	if fname, err = filepath.Abs(fname); err != nil {
		fmt.Println("Couldn't find Abs path of input.", err)
		os.Exit(1)
	}
	if fInfo, err = os.Stat(fname); err != nil {
		fmt.Println("Couldn't find source file to parse.", err)
		os.Exit(1)
	}

	WriteString(pkgCnt)
}

// WriteString is high level function that parses source
// and generates output code
func WriteString(pkgCnt string) {

	if err := p.Parse(fname, fInfo.IsDir()); err != nil {
		return
	}

	if out == "" {
		subpkg = ""
		if fInfo.IsDir() {
			out = filepath.Join(fname, p.PkgName+"_csvgen.go")
		} else {
			if s := strings.TrimSuffix(fname, ".go"); s == fname {
				// return errors.New("Filename must end in '.go'")
				fmt.Println("Filename must end in '.go'")
				os.Exit(1)
			} else {
				out = s + "_csvgen.go"
			}
		}
		pkgCnt = p.PkgName
	} else {
		if s := strings.TrimSuffix(out, ".go"); s == out {
			out = out + ".go"
		}
	}
	if subpkg != "" {
		if _, err := os.Stat(subpkg); os.IsNotExist(err) {
			fmt.Println(subpkg, "not exists. Trying to make directory")
			if err := os.Mkdir(subpkg, os.ModePerm); err != nil {
				fmt.Println("Couldn't make directory. Got an error:", err)
				os.Exit(1)
			}
		}
	}

	f = NewFile(pkgCnt)
	f.Comment("This code is generated by 'csvgen'")
	f.Comment("Do not edit")
	f.Line()

	GenerateCode()
	// fmt.Println(p.Error)
	// fmt.Println(p.StructMap)
	// fmt.Printf("%#v", f)
	if err := f.Save(filepath.Join(subpkg, out)); err != nil {
		fmt.Println("Couldn't save file. Got an error:", err)
		os.Exit(1)
	}
}

// GenerateCode generates code for every structure from input
func GenerateCode() {

	for _, v := range p.Structs {
		GenerateFuncs(v)
		f.Line()
	}

}

// GenerateFuncs processes every structure.
// Generated functions MarshalCSV and UnmarshalCSV process builtin types,
// call MarshalCSV and UnmarshalCSV for custom types
// and process pointer types assuming that pointers were initiated (memory is
// allocated).
// So the best way is to verify data structures before marshalling
// and unmarshalling for null pointers to prevent SEGFAULT
func GenerateFuncs(vstr parser.StructInfo) {

	// func (pv *Type) UnmarshalCSV(in []string) error {
	// 	if in == nil || len(in) < 2 {
	// 		return errors.New("Invalid input to UnmarshalCSV")
	// 	}
	//  unm_b, err := strconv.ParseBool(in[0])
	//  if err != nil {
	//      return err
	//  }
	//  pv.b = b
	//  unm_a, err := strconv.ParseInt(in[1], 10, 64)
	//  if err == nil {
	//      return err
	//  }
	//  this.a = a
	// }
	//
	// // func (this Type) MarshalCSV() []string {
	//      out := []string{}
	//      out = append(out, strconv.FormatInt(int64(this.a), 10))
	//      out = append(out, strconv.FormatBool(this.b))
	//      ...marshal logic
	//      return out, nil
	// }

	var unmarshallBody []Code
	var marshallBody []Code

	chkError := If(
		Id("in").Op("==").Id("nil").Op("||").Id("len").Call(Id("in")).Op("<").Lit(len(vstr.Fields)),
	).Block(
		Return().Qual("github.com/pkg/errors", "New").Call(Lit("Invalid input to *" + vstr.Name + " UnmarshalCSV")),
	)
	unmarshallBody = append(unmarshallBody, chkError)
	unmarshallBody = append(unmarshallBody, Id("i").Op(":=").Lit(0))
	marshallBody = append(marshallBody, Id("out").Op(":=").Index().String().Values())

	for ik, istr := range vstr.Fields {
		var g []Code
		var j *Statement
		star := ""
		ttype := istr.Type
		if istr.Type[0] == '*' {
			star = "*"
			ttype = istr.Type[1:]
		}

		unmarshallBody = append(unmarshallBody, nilCheck(star, istr.Name, istr.Type, false))
		marshallBody = append(marshallBody, nilCheck(star, istr.Name, istr.Type, true))

		switch ttype {
		case "bool":
			op := Qual("strconv", "ParseBool").Call(Id("in").Index(Id("i")))
			g = parseField(star, istr.Name, ttype, op, "false")
			j = marshalBody(Qual("strconv", "FormatBool").Call(Op(star).Id("pv").Op(".").Id(istr.Name)))
		case "float32":
			fallthrough
		case "float64":
			op := Qual("strconv", "ParseFloat").Call(List(Id("in").Index(Id("i")), Id(ttype[5:])))
			g = parseField(star, istr.Name, ttype, op, "0.0")
			j = marshalBody(Qual("strconv", "FormatFloat").
				Call(Op("float64").Call(Op(star).Id("pv").Op(".").Id(istr.Name)),
					LitRune('f'), Lit(-1), Id(ttype[5:])),
			)
		case "int":
			fallthrough
		case "int8":
			fallthrough
		case "int16":
			fallthrough
		case "int32":
			fallthrough
		case "int64":
			bn := ttype[3:]
			if bn == "" {
				bn = "0"
			}
			op := Qual("strconv", "ParseInt").Call(List(Id("in").Index(Id("i")), Lit(10), Id(bn)))
			g = parseField(star, istr.Name, ttype, op, "0")
			j = marshalBody(Qual("strconv", "FormatInt").
				Call(Op("int64").Call(Op(star).Id("pv").Op(".").Id(istr.Name)), Lit(10)),
			)
		case "uint":
			fallthrough
		case "uint8":
			fallthrough
		case "uint16":
			fallthrough
		case "uint32":
			fallthrough
		case "uint64":
			bn := ttype[4:]
			if bn == "" {
				bn = "0"
			}
			op := Qual("strconv", "ParseUint").Call(List(Id("in").Index(Id("i")), Lit(10), Id(bn)))
			g = parseField(star, istr.Name, ttype, op, "0")
			j = marshalBody(Qual("strconv", "FormatUint").
				Call(Op("uint64").Call(Op(star).Id("pv").Op(".").Id(istr.Name)), Lit(10)),
			)
		case "string":
			g = []Code{
				Op(star).Id("pv").Op(".").Id(istr.Name).Op("=").Id("in").Index(Id("i")),
			}
			j = marshalBody(Op(star).Id("pv").Op(".").Id(istr.Name))
		default:
			// By default generated code calls 'func (this *Type) UnmarshalCSV(s string) error'
			g = []Code{
				If(
					Err().Op(":=").Id("pv").Op(".").Id(istr.Name).Op(".").Id("UnmarshalCSV").
						Call(Id("in").Index(Id("i"))),
					Err().Op("!=").Nil(),
				).Block(
					Return().Err(),
				),
			}
			// By default generated code calls 'func (this Type) MarshalCSV() (string, error)'
			j = If(
				List(Id("mt"), Err()).Op(":=").Id("pv").Op(".").Id(istr.Name).Op(".").Id("MarshalCSV").
					Call(),
				Err().Op("!=").Nil(),
			).Block(
				Return(Id("out"), Err()),
			).Else().Block(
				marshalBody(Id("mt")),
			)
		}
		unmarshallBody = append(unmarshallBody, g...)
		marshallBody = append(marshallBody, j)
		if ik != (len(vstr.Fields) - 1) {
			unmarshallBody = append(unmarshallBody, Id("i").Op("++"))
		} else {
			unmarshallBody = append(unmarshallBody, Return().Id("nil"))
			marshallBody = append(marshallBody, Return().List(Id("out"), Id("nil")))
		}
	}

	if unmarshal {
		f.Comment("UnmarshalCSV " + vstr.Name + " func")
		f.Func().Params(
			Id("pv").Op("*").Id(vstr.Name),
		).Id("UnmarshalCSV").Params(
			Id("in").Index().String(),
		).Id("error").Block(
			unmarshallBody...,
		)

		f.Line()
	}

	if marshal {
		f.Comment("MarshalCSV " + vstr.Name + " func")
		f.Func().Params(
			Id("pv").Id(vstr.Name),
		).Id("MarshalCSV").Params().
			Parens(Index().String().Op(",").Id("error")).Block(
			marshallBody...,
		)
	}

	if unmarshal {
		AddList(vstr)
	}
}

func parseField(star string, fieldName string, fieldType string, op *Statement, defv string) []Code {
	fldNm := fieldPrefix + strings.Title(fieldName)
	parseBlock := []Code{
		List(Id(fldNm), Err()).Op(":=").Add(op),
		If(
			Err().Op("!=").Nil(),
		).Block(
			Return().Err(),
		),
	}

	var conv *Statement

	//  Check for float and int
	if fieldType[len(fieldType)-2:] != "64" {
		conv = Id(fieldType).Call(Id(fldNm))
	} else {
		conv = Id(fldNm)
	}

	parseBlock = append(parseBlock, Op(star).Id("pv").Op(".").Id(fieldName).Op("=").Add(conv))

	// Parse empty string with type rules
	// if in[i] == "" {
	// 	pv.Pp = 0
	// } else {
	//	...
	// }
	if parseEmpty {
		parseBlock = []Code{
			If(
				Id("in").Index(Id("i")).Op("==").Lit(""),
			).Block(
				Op(star).Id("pv").Op(".").Id(fieldName).Op("=").Id(defv),
			).Else().Block(
				parseBlock...,
			),
		}
	}
	return parseBlock
}

func marshalBody(typeRes *Statement) *Statement {
	return Id("out").Op("=").Append(Id("out"), Add(typeRes))
}

func nilCheck(star string, iname, itype string, marshall bool) *Statement {
	var s, t *Statement
	if marshall {
		t = List(Id("out"), Qual("github.com/pkg/errors", "New").Call(Lit("nil pointer found at "+iname+" "+itype)))
	} else {
		t = Qual("github.com/pkg/errors", "New").Call(Lit("nil pointer found at " + iname + " " + itype))
	}
	if star == "*" {
		s = If(
			Id("pv").Op(".").Id(iname).Op("==").Id("nil"),
		).Block(
			Return().Add(t),
		)
	} else {
		s = Null()
	}
	return s
}

func AddList(vstr parser.StructInfo) {
	// type FooList []Foo

	// func (pl *FooList) Push(in []string) error {
	// 	nf := Foo{}
	// 	if err := nf.UnmarshalCSV(in); err != nil {
	// 		return errors.Wrap(err, "Error in UnmarshalCSV Foo")
	// 	}
	// 	*pl = append(*pl, nf)
	// 	return nil
	// }

	f.Line()

	f.Comment(vstr.Name + "List csvparse.Pusher implementation")
	f.Type().Id(vstr.Name + "List").Index().Id(vstr.Name)

	f.Line()

	f.Comment("Push function for " + vstr.Name + "List struct")
	f.Func().Params(
		Id("pl").Op("*").Id(vstr.Name+"List"),
	).Id("Push").Params(Id("in").Index().String()).
		Parens(Error()).Block(
		Id("nf").Op(":=").Id(vstr.Name).Values(),
		If(
			Err().Op(":=").Id("nf").Op(".").Id("UnmarshalCSV").Call(Id("in")),
			Err().Op("!=").Nil(),
		).Block(
			Return(Qual("github.com/pkg/errors", "Wrapf").Call(
				Err(),
				Lit("Error in UnmarshalCSV "+vstr.Name+" in line %#v "),
				Id("in")),
			),
		),
		Op("*").Id("pl").Op("=").Append(Op("*").Id("pl"), Id("nf")),
		Return(Id("nil")),
	)

}