feat: add td.FlattenFilter function

Signed-off-by: Maxime Soulé <btik-git@scoubidou.com>

td/flatten.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020, Maxime Soulé
+// Copyright (c) 2020-2023, Maxime Soulé
 // All rights reserved.
 //
 // This source code is licensed under the BSD-style license found in the
@@ -8,9 +8,11 @@ package td
 
 import (
 	"reflect"
+	"strings"
 
 	"github.com/maxatome/go-testdeep/internal/color"
 	"github.com/maxatome/go-testdeep/internal/flat"
+	"github.com/maxatome/go-testdeep/internal/types"
 )
 
 // Flatten allows to flatten any slice, array or map in parameters of
@@ -84,6 +86,8 @@ import (
 //	td.Flatten(map[int]int{1: 2, 3: 4})
 //
 // is flattened as 1, 2, 3, 4 or 3, 4, 1, 2.
+//
+// See also [FlattenFilter].
 func Flatten(sliceOrMap any) flat.Slice {
 	switch reflect.ValueOf(sliceOrMap).Kind() {
 	case reflect.Slice, reflect.Array, reflect.Map:
@@ -92,3 +96,135 @@ func Flatten(sliceOrMap any) flat.Slice {
 		panic(color.BadUsage("Flatten(SLICE|ARRAY|MAP)", sliceOrMap, 1, true))
 	}
 }
+
+// FlattenFilter allows to flatten any slice, array or map in
+// parameters of operators expecting ...any after applying a function
+// on each item to exclude or transform it.
+//
+// fn must be a non-nil function with a signature like:
+//
+//	func(T) V
+//	func(T) (V, bool)
+//
+// T can be the same as V but it is not mandatory. The (V, bool)
+// returned case allows to exclude some items when returning false.
+//
+// If fn signature does not match these cases, FlattenFilter panics.
+//
+// If the type of an item of sliceOrMap is not convertible to T, the
+// item is dropped silently, as if fn returned false.
+//
+// fn can also be a string among:
+//
+//	"Smuggle:FIELD"
+//	"JSONPointer:/PATH"
+//
+// that are shortcuts for respectively:
+//
+//	func(in any) any { return td.Smuggle("FIELD", in) }
+//	func(in any) any { return td.JSONPointer("/PATH", in) }
+//
+// See [Smuggle] and [JSONPointer] for a description of what "FIELD"
+// and "/PATH" can really be.
+//
+// FlattenFilter can also be useful when testing some fields of
+// structs in a slice with [Set] or [Bag] families, here we test only
+// "Name" field:
+//
+//	type person struct {
+//	  Name string `json:"name"`
+//	  Age  int    `json:"age"`
+//	}
+//	got := []person{{"alice", 22}, {"bob", 18}, {"brian", 34}, {"britt", 32}}
+//
+//	td.Cmp(t, got,
+//		td.Bag(td.FlattenFilter(
+//			func(name string) any { return td.Smuggle("Name", name) },
+//			[]string{"alice", "britt", "brian", "bob"})))
+//
+//	td.Cmp(t, got,
+//		td.Bag(td.FlattenFilter(
+//			"Smuggle:Name", []string{"alice", "britt", "brian", "bob"})))
+//
+//	td.Cmp(t, got,
+//		td.Bag(td.FlattenFilter(
+//			func(name string) any { return td.JSONPointer("/name", name) },
+//			[]string{"alice", "britt", "brian", "bob"})))
+//
+//	td.Cmp(t, got,
+//		td.Bag(td.FlattenFilter(
+//			"JSONPointer:/name", []string{"alice", "britt", "brian", "bob"})))
+//
+//	td.Cmp(t, got,
+//		td.Bag(td.FlattenFilter(
+//			func(name string) any { return td.SuperJSONOf(`{"name":$1}`, name) },
+//			[]string{"alice", "britt", "brian", "bob"})))
+//
+//	td.Cmp(t, got,
+//		td.Bag(td.FlattenFilter(
+//			func(name string) any { return td.Struct(person{Name: name}) },
+//			[]string{"alice", "britt", "brian", "bob"})))
+//
+// See also [Flatten] and [Grep].
+func FlattenFilter(fn any, sliceOrMap any) flat.Slice {
+	const (
+		smugglePrefix     = "Smuggle:"
+		jsonPointerPrefix = "JSONPointer:"
+		usage             = "FlattenFilter(FUNC, SLICE|ARRAY|MAP)"
+		usageFunc         = usage + `, FUNC should be non-nil func(T) V or func(T) (V, bool) or a string "` + smugglePrefix + `…" or "` + jsonPointerPrefix + `…"`
+	)
+
+	// Smuggle & JSONPointer specific shortcuts
+	if s, ok := fn.(string); ok {
+		switch {
+		case strings.HasPrefix(s, smugglePrefix):
+			fn = func(in any) any {
+				return Smuggle(s[len(smugglePrefix):], in)
+			}
+		case strings.HasPrefix(s, jsonPointerPrefix):
+			fn = func(in any) any {
+				return JSONPointer(s[len(jsonPointerPrefix):], in)
+			}
+		default:
+			panic(color.Bad("usage: "+usageFunc+", but received %q as 1st parameter", s))
+		}
+	}
+
+	fnType := reflect.TypeOf(fn)
+	vfn := reflect.ValueOf(fn)
+
+	if fn == nil ||
+		fnType.Kind() != reflect.Func ||
+		fnType.NumIn() != 1 || fnType.IsVariadic() ||
+		(fnType.NumOut() != 1 && (fnType.NumOut() != 2 || fnType.Out(1) != types.Bool)) {
+		panic(color.BadUsage(usageFunc, fn, 1, false))
+	}
+	if vfn.IsNil() {
+		panic(color.Bad("usage: " + usageFunc))
+	}
+
+	inType := fnType.In(0)
+
+	switch reflect.ValueOf(sliceOrMap).Kind() {
+	case reflect.Slice, reflect.Array, reflect.Map:
+	default:
+		panic(color.BadUsage(usage, sliceOrMap, 2, true))
+	}
+
+	var final []any
+	for _, v := range flat.Values([]any{flat.Slice{Slice: sliceOrMap}}) {
+		if v.Type() != inType {
+			if !v.Type().ConvertibleTo(inType) {
+				continue
+			}
+			v = v.Convert(inType)
+		}
+
+		ret := vfn.Call([]reflect.Value{v})
+		if len(ret) == 1 || ret[1].Bool() {
+			final = append(final, ret[0].Interface())
+		}
+	}
+
+	return flat.Slice{Slice: final}
+}

td/flatten_test.go

@@ -7,6 +7,7 @@
 package td_test
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/maxatome/go-testdeep/internal/test"
@@ -33,3 +34,200 @@ func TestFlatten(t *testing.T) {
 	test.CheckPanic(t, func() { td.Flatten(42) },
 		"usage: Flatten(SLICE|ARRAY|MAP), but received int as 1st parameter")
 }
+
+func TestFlattenFilter(t *testing.T) {
+	t.Run("errors", func(t *testing.T) {
+		const usageFunc = `usage: FlattenFilter(FUNC, SLICE|ARRAY|MAP), FUNC should be non-nil func(T) V or func(T) (V, bool) or a string "Smuggle:…" or "JSONPointer:…"`
+		testCases := []struct {
+			name     string
+			fn       any
+			param    any
+			expected string
+		}{
+			{
+				name:     "untyped nil func",
+				param:    []int{},
+				expected: usageFunc + ", but received nil as 1st parameter",
+			},
+			{
+				name:     "not func",
+				fn:       42,
+				param:    []int{},
+				expected: usageFunc + ", but received int as 1st parameter",
+			},
+			{
+				name:     "func w/0 inputs",
+				fn:       func() int { return 0 },
+				param:    []int{},
+				expected: usageFunc + ", but received func() int as 1st parameter",
+			},
+			{
+				name:     "func w/2 inputs",
+				fn:       func(a, b int) int { return 0 },
+				param:    []int{},
+				expected: usageFunc + ", but received func(int, int) int as 1st parameter",
+			},
+			{
+				name:     "variadic func",
+				fn:       func(a ...int) int { return 0 },
+				param:    []int{},
+				expected: usageFunc + ", but received func(...int) int as 1st parameter",
+			},
+			{
+				name:     "func w/0 output",
+				fn:       func(a int) {},
+				param:    []int{},
+				expected: usageFunc + ", but received func(int) as 1st parameter",
+			},
+			{
+				name:     "func w/2 out without bool",
+				fn:       func(a int) (int, int) { return 0, 0 },
+				param:    []int{},
+				expected: usageFunc + ", but received func(int) (int, int) as 1st parameter",
+			},
+			{
+				name:     "bad shortcut",
+				fn:       "Pipo",
+				param:    []int{},
+				expected: usageFunc + `, but received "Pipo" as 1st parameter`,
+			},
+			{
+				name:     "typed nil func",
+				fn:       (func(a int) int)(nil),
+				param:    []int{},
+				expected: usageFunc,
+			},
+			{
+				name:     "untyped nil param",
+				fn:       func(a int) int { return 0 },
+				expected: "usage: FlattenFilter(FUNC, SLICE|ARRAY|MAP), but received nil as 2nd parameter",
+			},
+			{
+				name:     "untyped nil param",
+				fn:       func(a int) int { return 0 },
+				param:    42,
+				expected: "usage: FlattenFilter(FUNC, SLICE|ARRAY|MAP), but received int as 2nd parameter",
+			},
+		}
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				test.CheckPanic(t, func() { td.FlattenFilter(tc.fn, tc.param) },
+					tc.expected)
+			})
+		}
+	})
+
+	t.Run("ok", func(t *testing.T) {
+		cmp := func(t *testing.T, got, expected []any) {
+			t.Helper()
+
+			if (got == nil) != (expected == nil) {
+				t.Errorf("nil mismatch: got=%#v, expected=%#v", got, expected)
+				return
+			}
+
+			lg, le := len(got), len(expected)
+			l := lg
+			if l > le {
+				l = le
+			}
+			i := 0
+			for ; i < l; i++ {
+				if got[i] != expected[i] {
+					t.Errorf("#%d item differ, got=%v, expected=%v", i, got[i], expected[i])
+				}
+			}
+			for ; i < lg; i++ {
+				t.Errorf("#%d item is extra, got=%v", i, got[i])
+			}
+			for ; i < le; i++ {
+				t.Errorf("#%d item is missing, expected=%v", i, expected[i])
+			}
+		}
+
+		testCases := []struct {
+			name     string
+			fn       any
+			expected []any
+		}{
+			{
+				name:     "func never called",
+				fn:       func(s bool) bool { return true },
+				expected: nil,
+			},
+			{
+				name:     "double",
+				fn:       func(a int) int { return a * 2 },
+				expected: []any{0, 2, 4, 6, 8, 10, 12, 14, 16, 18},
+			},
+			{
+				name:     "even",
+				fn:       func(a int) (int, bool) { return a, a%2 == 0 },
+				expected: []any{0, 2, 4, 6, 8},
+			},
+			{
+				name:     "transform",
+				fn:       func(a int) (string, bool) { return strconv.Itoa(a), a%2 == 0 },
+				expected: []any{"0", "2", "4", "6", "8"},
+			},
+			{
+				name:     "nil",
+				fn:       func(a int) any { return nil },
+				expected: []any{nil, nil, nil, nil, nil, nil, nil, nil, nil, nil},
+			},
+			{
+				name: "convertible",
+				fn:   func(a int8) int8 { return a * 3 },
+				expected: []any{
+					int8(0), int8(3), int8(6), int8(9), int8(12),
+					int8(15), int8(18), int8(21), int8(24), int8(27),
+				},
+			},
+		}
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				s := td.FlattenFilter(tc.fn, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})
+				if sa, ok := s.Slice.([]any); test.IsTrue(t, ok) {
+					cmp(t, sa, tc.expected)
+				}
+			})
+		}
+	})
+
+	t.Run("complex", func(t *testing.T) {
+		type person struct {
+			Name string `json:"name"`
+			Age  int    `json:"age"`
+		}
+
+		got := []person{{"alice", 22}, {"bob", 18}, {"brian", 34}, {"britt", 32}}
+
+		td.Cmp(t, got,
+			td.Bag(td.FlattenFilter(
+				func(name string) any { return td.Smuggle("Name", name) },
+				[]string{"alice", "britt", "brian", "bob"})))
+
+		td.Cmp(t, got,
+			td.Bag(td.FlattenFilter(
+				"Smuggle:Name", []string{"alice", "britt", "brian", "bob"})))
+
+		td.Cmp(t, got,
+			td.Bag(td.FlattenFilter(
+				func(name string) any { return td.JSONPointer("/name", name) },
+				[]string{"alice", "britt", "brian", "bob"})))
+
+		td.Cmp(t, got,
+			td.Bag(td.FlattenFilter(
+				"JSONPointer:/name", []string{"alice", "britt", "brian", "bob"})))
+
+		td.Cmp(t, got,
+			td.Bag(td.FlattenFilter(
+				func(name string) any { return td.SuperJSONOf(`{"name":$1}`, name) },
+				[]string{"alice", "britt", "brian", "bob"})))
+
+		td.Cmp(t, got,
+			td.Bag(td.FlattenFilter(
+				func(name string) any { return td.Struct(person{Name: name}) },
+				[]string{"alice", "britt", "brian", "bob"})))
+	})
+}

td/td_grep.go

@@ -190,7 +190,7 @@ var _ TestDeep = &tdGrep{}
 //	td.Cmp(t, got, td.Grep(td.Gt(0), td.Nil()))     // succeeds
 //	td.Cmp(t, got, td.Grep(td.Gt(0), []int{}))      // fails
 //
-// See also [First] and [Last].
+// See also [First], [Last] and [FlattenFilter].
 func Grep(filter, expectedValue any) TestDeep {
 	g := tdGrep{}
 	g.initGrepBase(filter, expectedValue)

td/td_json_pointer.go

@@ -95,7 +95,7 @@ var _ TestDeep = &tdJSONPointer{}
 // TypeBehind method always returns nil as the expected type cannot be
 // guessed from a JSON pointer.
 //
-// See also [JSON], [SubJSONOf], [SuperJSONOf] and [Smuggle].
+// See also [JSON], [SubJSONOf], [SuperJSONOf], [Smuggle] and [FlattenFilter].
 //
 // [RFC 6901]: https://tools.ietf.org/html/rfc6901
 func JSONPointer(ptr string, expectedValue any) TestDeep {

td/td_smuggle.go

@@ -619,7 +619,7 @@ func buildCaster(outType reflect.Type, useString bool) reflect.Value {
 // fn. For the case where fn is a fields-path, it is always
 // any, as the type can not be known in advance.
 //
-// See also [Code] and [JSONPointer].
+// See also [Code], [JSONPointer] and [FlattenFilter].
 //
 // [json.RawMessage]: https://pkg.go.dev/encoding/json#RawMessage
 func Smuggle(fn, expectedValue any) TestDeep {