feat: Adds ExtractItems option to include key fields

Adds an option to `ExtractItems` to include key fields in the output.

Author: yongruilin

Diff for: typed/remove_test.go

@@ -885,3 +885,116 @@ func TestReversibleExtract(t *testing.T) {
 		})
 	}
 }
+
+type extractWithKeysTestCase struct {
+	name         string
+	rootTypeName string
+	schema       typed.YAMLObject
+	triplets     []extractTriplet
+}
+
+type extractTriplet struct {
+	object     typed.YAMLObject
+	set        *fieldpath.Set
+	wantOutput interface{}
+}
+
+var extractWithKeysCases = []extractWithKeysTestCase{{
+	name:         "associativeAndAtomicSchema",
+	rootTypeName: "myRoot",
+	schema:       typed.YAMLObject(associativeAndAtomicSchema),
+	triplets: []extractTriplet{
+		{
+			// extract with all key fields included
+			object: `{"list":[{"key":"nginx","id":1,"nv":2}]}`,
+			set: _NS(
+				_P("list", _KBF("key", "nginx", "id", 1), "key"),
+				_P("list", _KBF("key", "nginx", "id", 1), "id"),
+			),
+			wantOutput: typed.YAMLObject(`{"list":[{"key":"nginx","id":1}]}`),
+		},
+		{
+			// extract no key field included
+			object: `{"list":[{"key":"nginx","id":1,"nv":2}]}`,
+			set: _NS(
+				_P("list", _KBF("key", "nginx", "id", 1), "nv"),
+			),
+			wantOutput: typed.YAMLObject(`{"list":[{"key":"nginx","id":1, "nv":2}]}`),
+		},
+		{
+			// extract with patial keys included
+			object: `{"list":[{"key":"nginx","id":1,"nv":2}]}`,
+			set: _NS(
+				_P("list", _KBF("key", "nginx", "id", 1), "nv"),
+				_P("list", _KBF("key", "nginx", "id", 1), "id"),
+			),
+			wantOutput: typed.YAMLObject(`{"list":[{"key":"nginx","id":1, "nv":2}]}`),
+		},
+		{
+			// extract with null field value
+			object: `{"list":[{"key":"nginx","id":1,"nv":2}]}`,
+			set: _NS(
+				_P("list", _KBF("key", "nginx", "id", 1), "value"),
+			),
+			wantOutput: map[string]interface{}{
+				"list": []interface{}{nil},
+			},
+		},
+	},
+}}
+
+func (tt extractWithKeysTestCase) test(t *testing.T) {
+	parser, err := typed.NewParser(tt.schema)
+	if err != nil {
+		t.Fatalf("failed to create schema: %v", err)
+	}
+	pt := parser.Type(tt.rootTypeName)
+
+	for i, triplet := range tt.triplets {
+		triplet := triplet
+		t.Run(fmt.Sprintf("%v-valid-%v", tt.name, i), func(t *testing.T) {
+			t.Parallel()
+			// source typedValue obj
+			tv, err := pt.FromYAML(triplet.object)
+			if err != nil {
+				t.Fatal(err)
+			}
+			gotExtracted := tv.ExtractItems(triplet.set, typed.AppendKeyFields)
+
+			switch triplet.wantOutput.(type) {
+			case typed.YAMLObject:
+				wantOut, err := pt.FromYAML(triplet.wantOutput.(typed.YAMLObject))
+				if err != nil {
+					t.Fatalf("unable to parser/validate removeOutput yaml: %v\n%v", err, triplet.wantOutput)
+				}
+
+				if !value.Equals(gotExtracted.AsValue(), wantOut.AsValue()) {
+					t.Errorf("ExtractItems expected\n%v\nbut got\n%v\n",
+						value.ToString(wantOut.AsValue()), value.ToString(gotExtracted.AsValue()),
+					)
+				}
+			default:
+				// The extracted result
+				wantOut := value.NewValueInterface(triplet.wantOutput)
+				if !value.Equals(gotExtracted.AsValue(), wantOut) {
+					t.Errorf("ExtractItems expected\n%v\nbut got\n%v\n",
+						value.ToString(wantOut), value.ToString(gotExtracted.AsValue()),
+					)
+				}
+			}
+		})
+	}
+}
+
+// TestExtractWithKeys ensures that when you extract
+// items from an object with the AppendKeyField option,
+// the key fields are also included in the output.
+func TestExtractWithKeys(t *testing.T) {
+	for _, tt := range extractWithKeysCases {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			tt.test(t)
+		})
+	}
+}

Diff for: typed/typed.go

@@ -32,6 +32,14 @@ const (
 	AllowDuplicates ValidationOptions = iota
 )
 
+// ExtractItemsOptions is the list of all the options available when extracting items.
+type ExtractItemsOptions int
+
+const (
+	// AppendKeyFields means that when extracting items, the key field would also be included.
+	AppendKeyFields ExtractItemsOptions = iota
+)
+
 // AsTyped accepts a value and a type and returns a TypedValue. 'v' must have
 // type 'typeName' in the schema. An error is returned if the v doesn't conform
 // to the schema.
@@ -187,7 +195,33 @@ func (tv TypedValue) RemoveItems(items *fieldpath.Set) *TypedValue {
 }
 
 // ExtractItems returns a value with only the provided list or map items extracted from the value.
-func (tv TypedValue) ExtractItems(items *fieldpath.Set) *TypedValue {
+func (tv TypedValue) ExtractItems(items *fieldpath.Set, opts ...ExtractItemsOptions) *TypedValue {
+	for _, opt := range opts {
+		switch opt {
+		case AppendKeyFields:
+			tvPathSet, err := tv.ToFieldSet()
+			if err != nil {
+				continue
+			}
+			keyFieldPathSet := fieldpath.NewSet()
+			items.Iterate(func(path fieldpath.Path) {
+				if !tvPathSet.Has(path) {
+					return
+				}
+				for i, pe := range path {
+					if pe.Key == nil {
+						continue
+					}
+					for _, keyField := range *pe.Key {
+						keyName := keyField.Name
+						keyFieldPath := append(path[:i+1:i+1], fieldpath.PathElement{FieldName: &keyName})
+						keyFieldPathSet.Insert(keyFieldPath)
+					}
+				}
+			})
+			items = items.Union(keyFieldPathSet)
+		}
+	}
 	tv.value = removeItemsWithSchema(tv.value, items, tv.schema, tv.typeRef, true)
 	return &tv
 }