Add support for enums w/o needing to manually specify variations

Author: shane-tw

docparse/jsonschema.go

@@ -110,6 +110,7 @@ const (
 	paramOmitEmpty = "omitempty"
 	paramReadOnly  = "readonly"
 	paramOmitDoc   = "omitdoc"
+	paramEnum      = "enum"
 )
 
 func setTags(name, fName string, p *Schema, tags []string) error {
@@ -129,6 +130,9 @@ func setTags(name, fName string, p *Schema, tags []string) error {
 		case paramReadOnly:
 			t := true
 			p.Readonly = &t
+		case paramEnum:
+			// For this type of enum, we figure out the variations based on the type.
+			p.Type = "enum"
 
 		// Various string formats.
 		// https://tools.ietf.org/html/draft-handrews-json-schema-validation-01#section-7.3
@@ -274,6 +278,15 @@ start:
 
 	// Simple identifiers such as "string", "int", "MyType", etc.
 	case *ast.Ident:
+		if p.Type == "enum" && len(p.Enum) == 0 {
+			if variations, err := getEnumVariations(ref.File, pkg, typ.Name); len(variations) > 0 {
+				p.Enum = variations
+				return &p, nil
+			} else if err != nil {
+				return nil, err
+			}
+		}
+
 		mappedType, mappedFormat := MapType(prog, pkg+"."+typ.Name)
 		if mappedType == "" {
 			// Only check for canonicalType if this isn't mapped.
@@ -443,6 +456,33 @@ start:
 	return &p, nil
 }
 
+// Helper function to extract enum variations from a file.
+func getEnumVariations(currentFile, pkgPath, typeName string) ([]string, error) {
+	resolvedPath, pkg, err := resolvePackage(currentFile, pkgPath)
+	if err != nil {
+		return nil, fmt.Errorf("could not resolve package: %v", err)
+	}
+	decls, err := getDecls(pkg, resolvedPath)
+	if err != nil {
+		return nil, err
+	}
+	var variations []string
+	for _, decl := range decls {
+		if decl.vs == nil {
+			continue
+		}
+		if exprToString(decl.vs.Type) != typeName {
+			continue
+		}
+		if len(decl.vs.Values) == 0 {
+			continue
+		}
+		variations = append(variations, exprToString(decl.vs.Values[0]))
+	}
+
+	return variations, nil
+}
+
 func dropTypePointers(typ ast.Expr) ast.Expr {
 	var t *ast.StarExpr
 	var ok bool

docparse/jsonschema_test.go

@@ -22,6 +22,7 @@ func TestFieldToProperty(t *testing.T) {
 		"sliceP":    {Type: "array", Items: &Schema{Type: "string"}},
 		"cstr":      {Type: "string"},
 		"cstrP":     {Type: "string"},
+		"enumStr":   {Type: "enum", Enum: []string{"a", "b", "c"}},
 		"bar":       {Reference: "a.bar"},
 		"barP":      {Reference: "a.bar"},
 		"pkg":       {Reference: "mail.Address"},

docparse/testdata/src/a/a.go

@@ -11,17 +11,19 @@ import "net/mail"
 type foo struct {
 	// Documented str field.
 	// Newline.
-	str       string
-	byt       []byte
-	r         rune
-	b         bool // Inline docs.
-	fl        float64
-	err       error
-	strP      *string
-	slice     []string
-	sliceP    []*string
-	cstr      customStr
-	cstrP     *customStr
+	str    string
+	byt    []byte
+	r      rune
+	b      bool // Inline docs.
+	fl     float64
+	err    error
+	strP   *string
+	slice  []string
+	sliceP []*string
+	cstr   customStr
+	cstrP  *customStr
+	// {enum}
+	enumStr   customStr
 	bar       bar
 	barP      *bar
 	pkg       mail.Address
@@ -43,6 +45,12 @@ type nested struct {
 
 type customStr string
 
+const (
+	customStrA customStr = "a"
+	customStrB customStr = "b"
+	customStrC customStr = "c"
+)
+
 // Document me bar!
 type bar struct {
 	str string