feat: Adds ExtractItems option to include key fields

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

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 partial 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.WithAppendKeyFields())
+
+			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)
+		})
+	}
+}

typed/typed.go

@@ -32,6 +32,21 @@ const (
 	AllowDuplicates ValidationOptions = iota
 )
 
+// extractItemsOptions is the options available when extracting items.
+type extractItemsOptions struct {
+	appendKeyFields bool
+}
+
+type ExtractItemsOption func(*extractItemsOptions)
+
+// WithAppendKeyFields configures ExtractItems to include key fields.
+// It is exported for use in configuring ExtractItems.
+func WithAppendKeyFields() ExtractItemsOption {
+	return func(opts *extractItemsOptions) {
+		opts.appendKeyFields = true
+	}
+}
+
 // 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 +202,37 @@ 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 ...ExtractItemsOption) *TypedValue {
+	options := &extractItemsOptions{}
+	for _, opt := range opts {
+		opt(options)
+	}
+	if options.appendKeyFields {
+		tvPathSet, err := tv.ToFieldSet()
+		if err == nil {
+			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
+						// Create a new slice with the same elements as path[:i+1], but set its capacity to len(path[:i+1]).
+						// This ensures that appending to keyFieldPath creates a new underlying array, avoiding accidental
+						// modification of the original slice (path).
+						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
 }