Add docstring on valid types for list keys for list constructors (#466)

* Add docstring on valid types for list keys for list constructors

ygen/genstate.go

@@ -16,6 +16,7 @@ package ygen
 
 import (
 	"fmt"
+	"sort"
 	"strings"
 
 	"github.com/openconfig/goyang/pkg/yang"
@@ -66,6 +67,44 @@ func IsYgenDefinedGoType(t *MappedType) bool {
 	return t.IsEnumeratedValue || len(t.UnionTypes) >= 2 || t.NativeType == ygot.BinaryTypeName || t.NativeType == ygot.EmptyTypeName
 }
 
+// unionType is an internal type used to sort the UnionTypes map field of
+// MappedType. It satisfies sort.Interface.
+type unionType struct {
+	name  string
+	index int
+}
+
+type unionTypeList []unionType
+
+func (u unionTypeList) Len() int {
+	return len(u)
+}
+
+func (u unionTypeList) Swap(i, j int) {
+	u[i], u[j] = u[j], u[i]
+}
+
+func (u unionTypeList) Less(i, j int) bool {
+	return u[i].index < u[j].index
+}
+
+// OrderedUnionTypes returns a slice of union type names of the given
+// MappedType in YANG order. If the type is not a union (i.e. UnionTypes is
+// empty), then a nil slice is returned.
+func (t *MappedType) OrderedUnionTypes() []string {
+	var unionTypes unionTypeList
+	for name, index := range t.UnionTypes {
+		unionTypes = append(unionTypes, unionType{name: name, index: index})
+	}
+	sort.Sort(unionTypes)
+
+	var orderedUnionTypes []string
+	for _, unionType := range unionTypes {
+		orderedUnionTypes = append(orderedUnionTypes, unionType.name)
+	}
+	return orderedUnionTypes
+}
+
 // buildDirectoryDefinitions extracts the yang.Entry instances from a map of
 // entries that need struct or message definitions built for them. It resolves
 // each non-leaf yang.Entry to a Directory which contains the elements that are

ygen/genstate_test.go

@@ -24,6 +24,67 @@ import (
 	"github.com/openconfig/ygot/genutil"
 )
 
+func TestOrderedUnionTypes(t *testing.T) {
+	tests := []struct {
+		desc string
+		in   *MappedType
+		want []string
+	}{{
+		desc: "union type with 2 elements",
+		in: &MappedType{
+			NativeType: "A_Union",
+			UnionTypes: map[string]int{
+				"Binary":  1,
+				"float64": 2,
+			},
+		},
+		want: []string{
+			"Binary",
+			"float64",
+		},
+	}, {
+		desc: "union type with 3 elements",
+		in: &MappedType{
+			NativeType: "A_Union",
+			UnionTypes: map[string]int{
+				"uint64":  3,
+				"float64": 2,
+				"Binary":  1,
+			},
+		},
+		want: []string{
+			"Binary",
+			"float64",
+			"uint64",
+		},
+	}, {
+		desc: "non-union type",
+		in: &MappedType{
+			NativeType: "string",
+		},
+		want: nil,
+	}, {
+		desc: "union type with a single element",
+		in: &MappedType{
+			NativeType: "string",
+			UnionTypes: map[string]int{
+				"string": 0,
+			},
+		},
+		want: []string{
+			"string",
+		},
+	}}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			if diff := cmp.Diff(tt.in.OrderedUnionTypes(), tt.want); diff != "" {
+				t.Errorf("(-got, +want):\n%s", diff)
+			}
+		})
+	}
+}
+
 func TestBuildDirectoryDefinitions(t *testing.T) {
 	tests := []struct {
 		name                                    string

ypathgen/pathgen.go

@@ -460,6 +460,9 @@ type {{ .TypeName }}{{ .WildcardSuffix }} struct {
 	// path struct object.
 	goPathChildConstructorTemplate = mustTemplate("childConstructor", `
 // {{ .MethodName }} returns from {{ .Struct.TypeName }} the path struct for its child "{{ .SchemaName }}".
+{{- range $paramDocStr := .KeyParamDocStrs }}
+// {{ $paramDocStr }}
+{{- end }}
 func (n *{{ .Struct.TypeName }}) {{ .MethodName -}} ({{ .KeyParamListStr }}) *{{ .TypeName }} {
 	return &{{ .TypeName }}{
 		{{ .Struct.PathBaseTypeName }}: ygot.New{{ .Struct.PathBaseTypeName }}(
@@ -475,6 +478,7 @@ func (n *{{ .Struct.TypeName }}) {{ .MethodName -}} ({{ .KeyParamListStr }}) *{{
 	// builder style for the list API.
 	goKeyBuilderTemplate = mustTemplate("goKeyBuilder", `
 // {{ .MethodName }} sets {{ .TypeName }}'s key "{{ .KeySchemaName }}" to the specified value.
+// {{ .KeyParamDocStr }}
 func (n *{{ .TypeName }}) {{ .MethodName }}({{ .KeyParamName }} {{ .KeyParamType }}) *{{ .TypeName }} {
 	ygot.ModifyKey(n.NodePath, "{{ .KeySchemaName }}", {{ .KeyParamName }})
 	return n
@@ -640,6 +644,7 @@ type goPathFieldData struct {
 	Struct          goPathStructData // Struct stores template information for the field's containing struct.
 	KeyParamListStr string           // KeyParamListStr is the parameter list of the field's accessor method.
 	KeyEntriesStr   string           // KeyEntriesStr is an ordered list of comma-separated ("schemaName": unique camel-case name) for a list's keys.
+	KeyParamDocStrs []string         // KeyParamDocStrs is an ordered slice of docstrings documenting the types of each list key parameter.
 }
 
 // generateDirectorySnippet generates all Go code associated with a schema node
@@ -845,17 +850,19 @@ func generateChildConstructorsForListBuilderFormat(methodBuf *strings.Builder, l
 	for i := 0; i != keyN; i++ {
 		if err := goKeyBuilderTemplate.Execute(methodBuf,
 			struct {
-				MethodName    string
-				TypeName      string
-				KeySchemaName string
-				KeyParamType  string
-				KeyParamName  string
+				MethodName     string
+				TypeName       string
+				KeySchemaName  string
+				KeyParamType   string
+				KeyParamName   string
+				KeyParamDocStr string
 			}{
-				MethodName:    BuilderKeyPrefix + keyParams[i].varName,
-				TypeName:      fieldData.TypeName,
-				KeySchemaName: keyParams[i].name,
-				KeyParamName:  keyParams[i].varName,
-				KeyParamType:  keyParams[i].typeName,
+				MethodName:     BuilderKeyPrefix + keyParams[i].varName,
+				TypeName:       fieldData.TypeName,
+				KeySchemaName:  keyParams[i].name,
+				KeyParamName:   keyParams[i].varName,
+				KeyParamType:   keyParams[i].typeName,
+				KeyParamDocStr: keyParams[i].typeDocString,
 			}); err != nil {
 			errors = append(errors, err)
 		}
@@ -890,7 +897,7 @@ func generateChildConstructorsForList(methodBuf *strings.Builder, listAttr *ygen
 	// For each combination of parameter indices to be part of the method
 	// parameter list (i.e. NOT wildcarded).
 	for comboIndex, combo := range combos {
-		var paramListStrs, keyEntryStrs []string
+		var paramListStrs, paramDocStrs, keyEntryStrs []string
 		var anySuffixes []string
 
 		i := 0 // Loop through each parameter
@@ -903,15 +910,17 @@ func generateChildConstructorsForList(methodBuf *strings.Builder, listAttr *ygen
 			// Add selected parameters to the parameter list.
 			param := keyParams[paramIndex]
 			paramListStrs = append(paramListStrs, fmt.Sprintf("%s %s", param.varName, param.typeName))
+			paramDocStrs = append(paramDocStrs, param.typeDocString)
 			keyEntryStrs = append(keyEntryStrs, fmt.Sprintf(`"%s": %s`, param.name, param.varName))
 			i++
 		}
 		for ; i != keyN; i++ { // Handle edge case
 			keyEntryStrs = append(keyEntryStrs, fmt.Sprintf(`"%s": "*"`, keyParams[i].name))
 			anySuffixes = append(anySuffixes, WildcardSuffix+keyParams[i].varName)
 		}
-		// Create the string for the method parameter list and ygot.NodePath's key list.
+		// Create the string for the method parameter list, docstrings, and ygot.NodePath's key list.
 		fieldData.KeyParamListStr = strings.Join(paramListStrs, ", ")
+		fieldData.KeyParamDocStrs = paramDocStrs
 		fieldData.KeyEntriesStr = strings.Join(keyEntryStrs, ", ")
 
 		// Add wildcard description suffixes to the base method name
@@ -984,23 +993,26 @@ func getFieldTypeName(directory *ygen.Directory, directoryFieldName string, goFi
 }
 
 type keyParam struct {
-	name     string
-	varName  string
-	typeName string
+	name          string
+	varName       string
+	typeName      string
+	typeDocString string
 }
 
-// makeKeyParams generates the list of go parameter list components for a
-// child list's constructor method given the list's ygen.YangListAttr.
+// makeKeyParams generates the list of go parameter list components for a child
+// list's constructor method given the list's ygen.YangListAttr, as well as a
+// list of each parameter's types as a comment string.
 // It outputs the parameters in the same order as in the YangListAttr.
 // e.g.
 // in: &ygen.YangListAttr{
 // 	Keys: map[string]*ygen.MappedType{
 // 		"fluorine": &ygen.MappedType{NativeType: "string"},
-// 		"iodine-liquid":   &ygen.MappedType{NativeType: "Binary"},
+// 		"iodine-liquid":   &ygen.MappedType{NativeType: "A_Union", UnionTypes: {"Binary": 0, "uint64": 1}},
 // 	},
 // 	KeyElems: []*yang.Entry{{Name: "fluorine"}, {Name: "iodine-liquid"}},
 // }
-// out: [{"fluroine", "Fluorine", "string"}, {"iodine-liquid", "IodineLiquid", "oc.Binary"}]
+// param out: [{"fluroine", "Fluorine", "string"}, {"iodine-liquid", "IodineLiquid", "oc.A_Union"}]
+// docstring out: ["Fluorine: string", "IodineLiquid: [oc.Binary, oc.UnionUint64]"]
 func makeKeyParams(listAttr *ygen.YangListAttr, schemaStructPkgAccessor string) ([]keyParam, error) {
 	if len(listAttr.KeyElems) == 0 {
 		return nil, fmt.Errorf("makeKeyParams: invalid list - has no key; cannot process param list string")
@@ -1031,7 +1043,34 @@ func makeKeyParams(listAttr *ygen.YangListAttr, schemaStructPkgAccessor string)
 		default:
 			typeName = mappedType.NativeType
 		}
-		keyParams = append(keyParams, keyParam{name: keyElem.Name, varName: goKeyNameMap[keyElem.Name], typeName: typeName})
+		varName := goKeyNameMap[keyElem.Name]
+
+		typeDocString := typeName
+		if len(mappedType.UnionTypes) > 1 {
+			var genTypes []string
+			for _, name := range mappedType.OrderedUnionTypes() {
+				unionTypeName := name
+				if simpleName, ok := ygot.SimpleUnionBuiltinGoTypes[name]; ok {
+					unionTypeName = simpleName
+				}
+				// Add schemaStructPkgAccessor.
+				if strings.HasPrefix(unionTypeName, "*") {
+					unionTypeName = "*" + schemaStructPkgAccessor + unionTypeName[1:]
+				} else {
+					unionTypeName = schemaStructPkgAccessor + unionTypeName
+				}
+				genTypes = append(genTypes, unionTypeName)
+			}
+			// Create the subtype documentation string.
+			typeDocString = "[" + strings.Join(genTypes, ", ") + "]"
+		}
+
+		keyParams = append(keyParams, keyParam{
+			name:          keyElem.Name,
+			varName:       varName,
+			typeName:      typeName,
+			typeDocString: varName + ": " + typeDocString,
+		})
 	}
 	return keyParams, nil
 }