Add dyn.MapByPattern to map a function to values with matching paths

libs/dyn/pattern.go

@@ -0,0 +1,96 @@
+package dyn
+
+import (
+	"fmt"
+	"maps"
+	"slices"
+)
+
+// Pattern represents a matcher for paths in a [Value] configuration tree.
+// It is used by [MapByPattern] to apply a function to the values whose paths match the pattern.
+// Every [Path] is a valid [Pattern] that matches a single unique path.
+// The reverse is not true; not every [Pattern] is a valid [Path], as patterns may contain wildcards.
+type Pattern []patternComponent
+
+// A pattern component can visit a [Value] and recursively call into [visit] for matching elements.
+// Fixed components can match a single key or index, while wildcards can match any key or index.
+type patternComponent interface {
+	visit(v Value, prefix Path, suffix Pattern, opts visitOptions) (Value, error)
+}
+
+// NewPattern returns a new pattern from the given components.
+// The individual components may be created with [Key], [Index], or [Any].
+func NewPattern(cs ...patternComponent) Pattern {
+	return cs
+}
+
+// NewPatternFromPath returns a new pattern from the given path.
+func NewPatternFromPath(p Path) Pattern {
+	cs := make(Pattern, len(p))
+	for i, c := range p {
+		cs[i] = c
+	}
+	return cs
+}
+
+type anyKeyComponent struct{}
+
+// AnyKey returns a pattern component that matches any key.
+func AnyKey() patternComponent {
+	return anyKeyComponent{}
+}
+
+// This function implements the patternComponent interface.
+func (c anyKeyComponent) visit(v Value, prefix Path, suffix Pattern, opts visitOptions) (Value, error) {
+	m, ok := v.AsMap()
+	if !ok {
+		return InvalidValue, fmt.Errorf("expected a map at %q, found %s", prefix, v.Kind())
+	}
+
+	m = maps.Clone(m)
+	for key, value := range m {
+		var err error
+		nv, err := visit(value, prefix.Append(Key(key)), suffix, opts)
+		if err != nil {
+			// Leave the value intact if the suffix pattern didn't match any value.
+			if IsNoSuchKeyError(err) || IsIndexOutOfBoundsError(err) {
+				continue
+			}
+			return InvalidValue, err
+		}
+		m[key] = nv
+	}
+
+	return NewValue(m, v.Location()), nil
+}
+
+type anyIndexComponent struct{}
+
+// AnyIndex returns a pattern component that matches any index.
+func AnyIndex() patternComponent {
+	return anyIndexComponent{}
+}
+
+// This function implements the patternComponent interface.
+func (c anyIndexComponent) visit(v Value, prefix Path, suffix Pattern, opts visitOptions) (Value, error) {
+	s, ok := v.AsSequence()
+	if !ok {
+		return InvalidValue, fmt.Errorf("expected a sequence at %q, found %s", prefix, v.Kind())
+	}
+
+	s = slices.Clone(s)
+	for i, value := range s {
+		var err error
+		nv, err := visit(value, prefix.Append(Index(i)), suffix, opts)
+		if err != nil {
+			// Leave the value intact if the suffix pattern didn't match any value.
+			if IsNoSuchKeyError(err) || IsIndexOutOfBoundsError(err) {
+				continue
+			}
+			return InvalidValue, err
+		}
+		s[i] = nv
+	}
+
+	return NewValue(s, v.Location()), nil
+}

libs/dyn/pattern_test.go

@@ -0,0 +1,28 @@
+package dyn_test
+
+import (
+	"testing"
+
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewPattern(t *testing.T) {
+	pat := dyn.NewPattern(
+		dyn.Key("foo"),
+		dyn.Index(1),
+	)
+
+	assert.Len(t, pat, 2)
+}
+
+func TestNewPatternFromPath(t *testing.T) {
+	path := dyn.NewPath(
+		dyn.Key("foo"),
+		dyn.Index(1),
+	)
+
+	pat1 := dyn.NewPattern(dyn.Key("foo"), dyn.Index(1))
+	pat2 := dyn.NewPatternFromPath(path)
+	assert.Equal(t, pat1, pat2)
+}

libs/dyn/visit.go

@@ -47,7 +47,7 @@ type visitOptions struct {
 	fn func(Path, Value) (Value, error)
 }
 
-func visit(v Value, prefix, suffix Path, opts visitOptions) (Value, error) {
+func visit(v Value, prefix Path, suffix Pattern, opts visitOptions) (Value, error) {
 	if len(suffix) == 0 {
 		return opts.fn(prefix, v)
 	}
@@ -59,25 +59,31 @@ func visit(v Value, prefix, suffix Path, opts visitOptions) (Value, error) {
 	}
 
 	component := suffix[0]
-	prefix = prefix.Append(component)
 	suffix = suffix[1:]
 
+	// Visit the value with the current component.
+	return component.visit(v, prefix, suffix, opts)
+}
+
+func (component pathComponent) visit(v Value, prefix Path, suffix Pattern, opts visitOptions) (Value, error) {
+	path := prefix.Append(component)
+
 	switch {
 	case component.isKey():
 		// Expect a map to be set if this is a key.
 		m, ok := v.AsMap()
 		if !ok {
-			return InvalidValue, fmt.Errorf("expected a map to index %q, found %s", prefix, v.Kind())
+			return InvalidValue, fmt.Errorf("expected a map to index %q, found %s", path, v.Kind())
 		}
 
 		// Lookup current value in the map.
 		ev, ok := m[component.key]
 		if !ok {
-			return InvalidValue, noSuchKeyError{prefix}
+			return InvalidValue, noSuchKeyError{path}
 		}
 
 		// Recursively transform the value.
-		nv, err := visit(ev, prefix, suffix, opts)
+		nv, err := visit(ev, path, suffix, opts)
 		if err != nil {
 			return InvalidValue, err
 		}
@@ -100,17 +106,17 @@ func visit(v Value, prefix, suffix Path, opts visitOptions) (Value, error) {
 		// Expect a sequence to be set if this is an index.
 		s, ok := v.AsSequence()
 		if !ok {
-			return InvalidValue, fmt.Errorf("expected a sequence to index %q, found %s", prefix, v.Kind())
+			return InvalidValue, fmt.Errorf("expected a sequence to index %q, found %s", path, v.Kind())
 		}
 
 		// Lookup current value in the sequence.
 		if component.index < 0 || component.index >= len(s) {
-			return InvalidValue, indexOutOfBoundsError{prefix}
+			return InvalidValue, indexOutOfBoundsError{path}
 		}
 
 		// Recursively transform the value.
 		ev := s[component.index]
-		nv, err := visit(ev, prefix, suffix, opts)
+		nv, err := visit(ev, path, suffix, opts)
 		if err != nil {
 			return InvalidValue, err
 		}

libs/dyn/visit_get.go

@@ -14,7 +14,7 @@ func Get(v Value, path string) (Value, error) {
 // If the path doesn't exist, it returns InvalidValue and an error.
 func GetByPath(v Value, p Path) (Value, error) {
 	out := InvalidValue
-	_, err := visit(v, EmptyPath, p, visitOptions{
+	_, err := visit(v, EmptyPath, NewPatternFromPath(p), visitOptions{
 		fn: func(_ Path, ev Value) (Value, error) {
 			// Capture the value argument to return it.
 			out = ev

libs/dyn/visit_map.go

@@ -40,7 +40,7 @@ func Foreach(fn MapFunc) MapFunc {
 	}
 }
 
-// Map applies the given function to the value at the specified path in the specified value.
+// Map applies a function to the value at the given path in the given value.
 // It is identical to [MapByPath], except that it takes a string path instead of a [Path].
 func Map(v Value, path string, fn MapFunc) (Value, error) {
 	p, err := NewPathFromString(path)
@@ -50,15 +50,21 @@ func Map(v Value, path string, fn MapFunc) (Value, error) {
 	return MapByPath(v, p, fn)
 }
 
-// Map applies the given function to the value at the specified path in the specified value.
+// MapByPath applies a function to the value at the given path in the given value.
+// It is identical to [MapByPattern], except that it takes a [Path] instead of a [Pattern].
+// This means it only matches a single value, not a pattern of values.
+func MapByPath(v Value, p Path, fn MapFunc) (Value, error) {
+	return MapByPattern(v, NewPatternFromPath(p), fn)
+}
+
+// MapByPattern applies a function to the values whose paths match the given pattern in the given value.
 // If successful, it returns the new value with all intermediate values copied and updated.
 //
-// If the path contains a key that doesn't exist, or an index that is out of bounds,
-// it returns the original value and no error. This is because setting a value at a path
-// that doesn't exist is a no-op.
+// If the pattern contains a key that doesn't exist, or an index that is out of bounds,
+// it returns the original value and no error.
 //
-// If the path is invalid for the given value, it returns InvalidValue and an error.
-func MapByPath(v Value, p Path, fn MapFunc) (Value, error) {
+// If the pattern is invalid for the given value, it returns InvalidValue and an error.
+func MapByPattern(v Value, p Pattern, fn MapFunc) (Value, error) {
 	nv, err := visit(v, EmptyPath, p, visitOptions{
 		fn: fn,
 	})

libs/dyn/visit_map_test.go

@@ -268,3 +268,107 @@ func TestMapForeachOnOtherError(t *testing.T) {
 	}))
 	assert.ErrorContains(t, err, "expected a map or sequence, found int")
 }
+
+func TestMapByPatternOnNilValue(t *testing.T) {
+	var err error
+	_, err = dyn.MapByPattern(dyn.NilValue, dyn.NewPattern(dyn.AnyKey()), nil)
+	assert.ErrorContains(t, err, `expected a map at "", found nil`)
+	_, err = dyn.MapByPattern(dyn.NilValue, dyn.NewPattern(dyn.AnyIndex()), nil)
+	assert.ErrorContains(t, err, `expected a sequence at "", found nil`)
+}
+
+func TestMapByPatternOnMap(t *testing.T) {
+	vin := dyn.V(map[string]dyn.Value{
+		"a": dyn.V(map[string]dyn.Value{
+			"b": dyn.V(42),
+		}),
+		"b": dyn.V(map[string]dyn.Value{
+			"c": dyn.V(43),
+		}),
+	})
+
+	var err error
+
+	// Expect an error if the pattern structure doesn't match the value structure.
+	_, err = dyn.MapByPattern(vin, dyn.NewPattern(dyn.AnyKey(), dyn.Index(0)), nil)
+	assert.ErrorContains(t, err, `expected a sequence to index`)
+
+	// Apply function to pattern "*.b".
+	vout, err := dyn.MapByPattern(vin, dyn.NewPattern(dyn.AnyKey(), dyn.Key("b")), func(p dyn.Path, v dyn.Value) (dyn.Value, error) {
+		assert.Equal(t, dyn.NewPath(dyn.Key("a"), dyn.Key("b")), p)
+		assert.Equal(t, dyn.V(42), v)
+		return dyn.V(44), nil
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, map[string]any{
+		"a": map[string]any{
+			"b": 44,
+		},
+		"b": map[string]any{
+			"c": 43,
+		},
+	}, vout.AsAny())
+}
+
+func TestMapByPatternOnMapWithoutMatch(t *testing.T) {
+	vin := dyn.V(map[string]dyn.Value{
+		"a": dyn.V(map[string]dyn.Value{
+			"b": dyn.V(42),
+		}),
+		"b": dyn.V(map[string]dyn.Value{
+			"c": dyn.V(43),
+		}),
+	})
+
+	// Apply function to pattern "*.zzz".
+	vout, err := dyn.MapByPattern(vin, dyn.NewPattern(dyn.AnyKey(), dyn.Key("zzz")), nil)
+	assert.NoError(t, err)
+	assert.Equal(t, vin, vout)
+}
+
+func TestMapByPatternOnSequence(t *testing.T) {
+	vin := dyn.V([]dyn.Value{
+		dyn.V([]dyn.Value{
+			dyn.V(42),
+		}),
+		dyn.V([]dyn.Value{
+			dyn.V(43),
+			dyn.V(44),
+		}),
+	})
+
+	var err error
+
+	// Expect an error if the pattern structure doesn't match the value structure.
+	_, err = dyn.MapByPattern(vin, dyn.NewPattern(dyn.AnyIndex(), dyn.Key("a")), nil)
+	assert.ErrorContains(t, err, `expected a map to index`)
+
+	// Apply function to pattern "*.c".
+	vout, err := dyn.MapByPattern(vin, dyn.NewPattern(dyn.AnyIndex(), dyn.Index(1)), func(p dyn.Path, v dyn.Value) (dyn.Value, error) {
+		assert.Equal(t, dyn.NewPath(dyn.Index(1), dyn.Index(1)), p)
+		assert.Equal(t, dyn.V(44), v)
+		return dyn.V(45), nil
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, []any{
+		[]any{42},
+		[]any{43, 45},
+	}, vout.AsAny())
+}
+
+func TestMapByPatternOnSequenceWithoutMatch(t *testing.T) {
+	vin := dyn.V([]dyn.Value{
+		dyn.V([]dyn.Value{
+			dyn.V(42),
+		}),
+		dyn.V([]dyn.Value{
+			dyn.V(43),
+			dyn.V(44),
+		}),
+	})
+
+	// Apply function to pattern "*.zzz".
+	vout, err := dyn.MapByPattern(vin, dyn.NewPattern(dyn.AnyIndex(), dyn.Index(42)), nil)
+	assert.NoError(t, err)
+	assert.Equal(t, vin, vout)
+}

libs/dyn/visit_set.go

@@ -25,10 +25,10 @@ func SetByPath(v Value, p Path, nv Value) (Value, error) {
 		return nv, nil
 	}
 
-	parent := p[:lp-1]
 	component := p[lp-1]
+	p = p[:lp-1]
 
-	return visit(v, EmptyPath, parent, visitOptions{
+	return visit(v, EmptyPath, NewPatternFromPath(p), visitOptions{
 		fn: func(prefix Path, v Value) (Value, error) {
 			path := prefix.Append(component)