Merge pull request #4 from aaharu/enum-tmp

enum

README.md

@@ -28,10 +28,78 @@ OPTIONS
            Show version.
 ```
 
+## Examples
+
+```bash
+schemarshal -p sample test_data/a.json
+```
+```go
+// generated by schemarshal 0.3.0 `schemarshal -p sample test_data/a.json`
+
+package sample
+
+import (
+	"fmt"
+	"time"
+)
+
+type A struct {
+	PhoneNumber []struct {
+		Location string `json:"location"`
+		Code     int    `json:"code"`
+		Hoge     struct {
+			Aaa bool `json:"aaa,omitempty"`
+		} `json:"hoge,omitempty"`
+		Fuga map[string]interface{} `json:"fuga,omitempty"`
+	} `json:"phoneNumber"`
+	Address struct {
+		City          *string `json:"city"`
+		StreetAddress string  `json:"streetAddress"`
+	} `json:"address"`
+	Test  []int     `json:"test,omitempty"`
+	Test2 time.Time `json:"test2,omitempty"`
+}
+
+```
+
+```bash
+curl -s "https://raw.githubusercontent.com/aaharu/schemarshal/master/test_data/disk.json" | schemarshal
+```
+```go
+// generated by schemarshal 0.3.0 `schemarshal`
+
+package main
+
+import (
+	"fmt"
+	"time"
+)
+
+type Disk struct {
+	Device string `json:"device"`
+	Type   string `json:"type"`
+}
+
+type TypeEnum int
+
+const (
+	enumDisk TypeEnum = iota
+	enumDisk2
+)
+
+func (enum TypeEnum) MarshalJSON() ([]byte, error) {
+	var varType = []interface{}{"disk", "disk2"}
+	return []byte(fmt.Sprintf("%v", varType[enum])), nil
+}
+
+```
+
 ## TODO
 
-- [ ] support enum
-- [ ] unit tests
+- [x] support enum
+- [x] unit tests
+- [ ] enum's bug fix
+- [ ] add test cases
 - [ ] write doc
 
 ## Dependencies

codegen/generator.go

@@ -8,10 +8,14 @@ import (
 	"bytes"
 	"fmt"
 	"go/format"
+	"strings"
 
+	"github.com/aaharu/schemarshal/utils"
 	"github.com/aaharu/schemarshal/version"
 )
 
+var enumList = map[string][]interface{}{} // FIXME
+
 // Generator of Go source code from JSON Schema
 type Generator struct {
 	name    string // package nage
@@ -29,12 +33,12 @@ func NewGenerator(packageName string, command string) *Generator {
 }
 
 // AddImport add an import statement
-func (g *Generator) AddImport(path string, name *string) {
+func (g *Generator) AddImport(path string, name string) {
 	if g.imports == nil {
 		g.imports = []*importSpec{}
 	}
 	g.imports = append(g.imports, &importSpec{
-		name: *name,
+		name: name,
 		path: path,
 	})
 }
@@ -76,6 +80,41 @@ func (g *Generator) Generate() ([]byte, error) {
 		}
 	}
 
+	if len(enumList) > 0 {
+		buf.WriteString("\n")
+		for typeName, enum := range enumList {
+			buf.WriteString("type " + typeName + "Enum int\n")
+			buf.WriteString("const (\n")
+			for i := range enum {
+				if i == 0 {
+					buf.WriteString("enum" + utils.Ucfirst(fmt.Sprintf("%v", enum[0])) + " " + typeName + "Enum = iota\n")
+				} else {
+					buf.WriteString("enum" + utils.Ucfirst(fmt.Sprintf("%v", enum[i])))
+				}
+			}
+			buf.WriteString(")\n")
+			buf.WriteString("func (enum " + typeName + "Enum) MarshalJSON() ([]byte, error) {\n")
+
+			buf.WriteString("var var" + typeName + " = []interface{}{")
+			for i := range enum {
+				switch v := enum[i].(type) {
+				case string:
+					buf.WriteString(`"`)
+					buf.WriteString(strings.Replace(v, `"`, `\"`, -1))
+					buf.WriteString(`"`)
+				default:
+					buf.WriteString(fmt.Sprintf("%v", v))
+				}
+				buf.WriteString(",")
+			}
+			buf.WriteString("}\n")
+
+			buf.WriteString("return []byte(fmt.Sprintf(\"%v\", var" + typeName + "[enum])), nil")
+
+			buf.WriteString("}\n\n")
+		}
+	}
+
 	return format.Source(buf.Bytes())
 }
 
@@ -89,25 +128,28 @@ type typeSpec struct {
 	jsontype *JSONType
 }
 
-type jsonformat int
+type jsonFormat int
 
 const (
-	OBJECT jsonformat = iota
-	ARRAY
-	STRING
-	BOOLEAN
-	NUMBER
-	INTEGER
-	DATETIME
+	formatObject jsonFormat = iota
+	formatArray
+	formatString
+	formatBoolean
+	formatNumber
+	formatInteger
+	formatDatetime
 )
 
+// JSONType ...
 type JSONType struct {
-	format   jsonformat
+	format   jsonFormat
 	nullable bool
-	fields   []*field // object
-	itemType *JSONType
+	fields   []*field  // only object
+	itemType *JSONType // only array
+	enumType string
 }
 
+// AddField ...
 func (t *JSONType) AddField(f *field) {
 	if t.fields == nil {
 		t.fields = []*field{}
@@ -120,7 +162,7 @@ func (t *JSONType) generate() []byte {
 	if t.nullable {
 		buf.WriteString("*")
 	}
-	if t.format == OBJECT {
+	if t.format == formatObject {
 		if t.fields == nil {
 			buf.WriteString("map[string]interface{}")
 		} else {
@@ -135,18 +177,18 @@ func (t *JSONType) generate() []byte {
 			}
 			buf.WriteString("}")
 		}
-	} else if t.format == ARRAY {
+	} else if t.format == formatArray {
 		buf.WriteString("[]")
 		buf.Write(t.itemType.generate())
-	} else if t.format == STRING {
+	} else if t.format == formatString {
 		buf.WriteString("string")
-	} else if t.format == BOOLEAN {
+	} else if t.format == formatBoolean {
 		buf.WriteString("bool")
-	} else if t.format == NUMBER {
+	} else if t.format == formatNumber {
 		buf.WriteString("float64")
-	} else if t.format == INTEGER {
+	} else if t.format == formatInteger {
 		buf.WriteString("int")
-	} else if t.format == DATETIME {
+	} else if t.format == formatDatetime {
 		buf.WriteString("time.Time")
 	}
 	return buf.Bytes()

codegen/jsonschema.go

@@ -22,8 +22,8 @@ type JSONSchema struct {
 	schema *schema.Schema
 }
 
-// Read and initialize struct
-func Read(input io.Reader) (*JSONSchema, error) {
+// ReadSchema and initialize struct
+func ReadSchema(input io.Reader) (*JSONSchema, error) {
 	schema, err := schema.Read(input)
 	if err != nil {
 		return nil, err
@@ -47,38 +47,50 @@ func (js *JSONSchema) GetTitle() string {
 }
 
 // Parse return JSON Schema type
-func (js *JSONSchema) Parse() (*JSONType, error) {
-	if inPrimitiveTypes(schema.IntegerType, js.schema.Type) {
-		t := &JSONType{
-			format: INTEGER,
+func (js *JSONSchema) Parse(fieldName string) (*JSONType, error) {
+	t := &JSONType{}
+	if inPrimitiveTypes(schema.IntegerType, js.schema.Type) ||
+		inPrimitiveTypes(schema.BooleanType, js.schema.Type) ||
+		inPrimitiveTypes(schema.NumberType, js.schema.Type) {
+		if inPrimitiveTypes(schema.IntegerType, js.schema.Type) {
+			t.format = formatInteger
+		} else if inPrimitiveTypes(schema.BooleanType, js.schema.Type) {
+			t.format = formatBoolean
+		} else {
+			t.format = formatNumber
 		}
 		if inPrimitiveTypes(schema.NullType, js.schema.Type) {
 			t.nullable = true
 		}
+		if js.schema.Enum != nil {
+			enumList[fieldName] = js.schema.Enum
+		}
+		t.enumType = fieldName
 		return t, nil
 	}
 	if inPrimitiveTypes(schema.StringType, js.schema.Type) {
-		t := &JSONType{}
 		if js.schema.Format == schema.FormatDateTime {
-			t.format = DATETIME
+			t.format = formatDatetime
 		} else {
-			t.format = STRING
+			t.format = formatString
 		}
 		if inPrimitiveTypes(schema.NullType, js.schema.Type) {
 			t.nullable = true
 		}
+		if js.schema.Enum != nil {
+			enumList[fieldName] = js.schema.Enum
+		}
+		t.enumType = fieldName
 		return t, nil
 	}
 	if inPrimitiveTypes(schema.ObjectType, js.schema.Type) {
-		t := &JSONType{
-			format: OBJECT,
-		}
+		t.format = formatObject
 		if inPrimitiveTypes(schema.NullType, js.schema.Type) {
 			t.nullable = true
 		}
 		if js.schema.Properties != nil {
 			for key, propSchema := range js.schema.Properties {
-				propType, err := NewSchema(propSchema).Parse()
+				propType, err := NewSchema(propSchema).Parse(utils.Ucfirst(key))
 				if err != nil {
 					return nil, err
 				}
@@ -92,47 +104,31 @@ func (js *JSONSchema) Parse() (*JSONType, error) {
 				})
 			}
 		}
+		if js.schema.Enum != nil {
+			enumList[fieldName] = js.schema.Enum
+		}
+		t.enumType = fieldName
 		return t, nil
 	}
 	if inPrimitiveTypes(schema.ArrayType, js.schema.Type) {
 		if js.schema.Items.TupleMode {
 			// unsupported
 			err := fmt.Errorf("unsupported type %v", js.schema.Items)
-			return nil, err
-		}
-		t := &JSONType{
-			format: ARRAY,
+			return t, err
 		}
+		t.format = formatArray
 		if inPrimitiveTypes(schema.NullType, js.schema.Type) {
 			t.nullable = true
 		}
-		itemType, err := NewSchema(js.schema.Items.Schemas[0]).Parse()
+		itemType, err := NewSchema(js.schema.Items.Schemas[0]).Parse("")
 		if err != nil {
 			return nil, err
 		}
 		t.itemType = itemType
 		return t, nil
 	}
-	if inPrimitiveTypes(schema.BooleanType, js.schema.Type) {
-		t := &JSONType{
-			format: BOOLEAN,
-		}
-		if inPrimitiveTypes(schema.NullType, js.schema.Type) {
-			t.nullable = true
-		}
-		return t, nil
-	}
-	if inPrimitiveTypes(schema.NumberType, js.schema.Type) {
-		t := &JSONType{
-			format: NUMBER,
-		}
-		if inPrimitiveTypes(schema.NullType, js.schema.Type) {
-			t.nullable = true
-		}
-		return t, nil
-	}
 	err := fmt.Errorf("unsupported type %v", js.schema.Type)
-	return nil, err
+	return t, err
 }
 
 func inPrimitiveTypes(needle schema.PrimitiveType, haystack schema.PrimitiveTypes) bool {

codegen/jsonschema_test.go

@@ -7,13 +7,26 @@ package codegen
 import (
 	"os"
 	"testing"
+
+	"github.com/aaharu/schemarshal/utils"
 )
 
-func Test(t *testing.T) {
+func TestSample1(t *testing.T) {
 	file, _ := os.Open("../test_data/a.json")
 	defer file.Close()
-	js, _ := Read(file)
-	jsType, _ := js.Parse()
+	js, _ := ReadSchema(file)
+	jsType, _ := js.Parse(utils.Ucfirst(utils.FileName(file)))
+	actual := jsType.generate()
+	if len(actual) < 1 {
+		t.Errorf("got %v\n", string(actual))
+	}
+}
+
+func TestSample2(t *testing.T) {
+	file, _ := os.Open("../test_data/disk.json")
+	defer file.Close()
+	js, _ := ReadSchema(file)
+	jsType, _ := js.Parse(utils.Ucfirst(utils.FileName(file)))
 	actual := jsType.generate()
 	if len(actual) < 1 {
 		t.Errorf("got %v\n", string(actual))