generator.go

package structgen

import (
	"bytes"
	"io"
	"strings"
)

type Generator struct {
	packageName string
	refMap      SchemaMap
	rootName    string
	schema      *Schema
	structs     StructMap
}

func NewGenerator(root, packageName string, schema *Schema) *Generator {
	return &Generator{
		packageName: packageName,
		rootName:    root,
		schema:      schema,
	}
}

func (g *Generator) Read(p []byte) (int, error) {
	if g.structs == nil {
		g.structs = make(StructMap)
		g.refMap = make(SchemaMap)
		// build up ref map first
		for d, s := range g.schema.Definitions {
			if s.ID != "" || d != "" {
				if s.ID != "" {
					g.refMap[s.ID] = s
				} else { // some more generic types have missing ids, prefix them
					g.refMap["#/definitions/"+d] = s
				}
			}
		}
		g.resolve(g.rootName, g.schema)
	}

	buf := &bytes.Buffer{}
	_, err := buf.WriteString(strings.Replace(headerTemplate, "{{ name }}", g.packageName, 1))
	if err != nil {
		return 0, err
	}

	for _, s := range g.structs {
		ss, err := s.StructString()
		if err != nil {
			return 0, err
		}
		_, err = buf.WriteString(ss)
		if err != nil {
			return 0, err
		}
	}
	n, err := buf.Read(p)
	if err != nil {
		return n, err
	}
	return n, io.EOF
}

func (g *Generator) resolve(fieldName string, schema *Schema) {
	structName := schema.toCamelCase(fieldName)
	if _, ok := g.structs[structName]; ok {
		return
	}
	for n, s := range schema.Properties {
		g.resolve(n, s)
	}
	for _, s := range schema.PatternProperties {
		g.resolve("", s)
	}
	for d, s := range schema.Definitions {
		g.resolve(d, s)
	}
	g.structs[structName] = g.newStruct(fieldName, schema)
}

func (g *Generator) newStruct(name string, s *Schema) *Struct {
	fields := make(map[string]*Struct)
	required := func(str string) bool {
		for _, r := range s.Required {
			if str == r {
				return true
			}
		}
		return false
	}
	item := s
	if t, _ := item.GetType(); t == TypeArray && item.Items != nil {
		item = item.Items
	}
	for n, prop := range item.Properties {
		propStruct := g.newStruct(n, prop)
		propStruct.Required = required(n)
		fields[n] = propStruct

		if t, _ := prop.GetType(); !(t == TypeObject || t == TypeArray) {
			continue
		}

		if _, ok := g.structs[propStruct.Name]; ok {
			println("Err:", propStruct.Name, "already exists")
		} else {
			g.structs[propStruct.Name] = propStruct
		}
	}
	t := s.getGoType(name)
	if refName, ref := g.getReference(s.Reference); ref != nil {
		t = ref.getGoType(refName)
		cc := ref.toCamelCase(refName)
		if _, ok := g.structs[cc]; !ok {
			g.structs[cc] = g.newStruct(refName, ref)
		}
	}
	return &Struct{
		Comment:  g.createComment(s),
		Fields:   fields,
		ID:       s.ID,
		JSONName: name,
		Name:     s.toCamelCase(name),
		Type:     t,
	}
}

func (g *Generator) getReference(id string) (string, *Schema) {
	if id == "" {
		return "", nil
	}
	for n, s := range g.schema.Definitions {
		if s.ID == id {
			return n, s
		}
	}
	return "", nil
}

// createComment writes the type options for interface types to a comment string.
func (g *Generator) createComment(s *Schema) (comment string) {
	var list []*Schema
	if t, types := s.GetType(); len(types) > 1 {
		comment = strings.Join(types, ",")
	} else if t == TypeArray {
		if _, atypes := s.Items.GetType(); len(atypes) > 1 {
			comment = strings.Join(atypes, ",")
		} else {
			list = s.Items.OneOf
		}
	} else if len(s.OneOf) > 0 {
		list = s.OneOf
	}
	if len(list) > 0 {
		var types []string
		for _, s := range list {
			_, st := s.GetType()
			types = append(types, st...)
		}
		comment = strings.Join(types, ",")
	}
	return
}

var headerTemplate = `// Code generated by 'github.com/sevenval/structgen'. DO NOT EDIT.

package {{ name }}

`