Merge pull request #61 from zoncoen/fix

fix: apply custom extract funcs even if the field is inline
zoncoen · Jan 25, 2024 · ed40571 · ed40571
2 parents 20ae25e + c42c50d
commit ed40571
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 27 deletions.
diff --git a/key.go b/key.go
@@ -24,11 +24,12 @@ type KeyExtractorContext interface {
 
 // Key represents an extractor to access the value by key.
 type Key struct {
-	key             string
-	caseInsensitive bool
-	structTags      []string
-	fieldNameGetter func(f reflect.StructField) string
-	isInlineFuncs   []func(reflect.StructField) bool
+	key                string
+	caseInsensitive    bool
+	structTags         []string
+	customExtractFuncs []func(ExtractFunc) ExtractFunc
+	fieldNameGetter    func(f reflect.StructField) string
+	isInlineFuncs      []func(reflect.StructField) bool
 }
 
 // Extract extracts the value from v by key.
@@ -116,13 +117,19 @@ func (e *Key) extract(v reflect.Value) (reflect.Value, bool) {
 				inlines = append(inlines, i)
 			}
 		}
-		for _, i := range inlines {
-			val, ok := e.extract(v.FieldByIndex([]int{i}))
-			if ok {
-				if isUnexportedField(val) {
-					unexported = &val
-				} else {
-					return val, true
+		if len(inlines) > 0 {
+			f := e.Extract
+			for i := len(e.customExtractFuncs) - 1; i >= 0; i-- {
+				f = e.customExtractFuncs[i](f)
+			}
+			for _, i := range inlines {
+				val, ok := f(v.FieldByIndex([]int{i}))
+				if ok {
+					if isUnexportedField(val) {
+						unexported = &val
+					} else {
+						return val, true
+					}
 				}
 			}
 		}

diff --git a/key_test.go b/key_test.go
@@ -59,12 +59,13 @@ type AnonymousField struct {
 func TestKey_Extract(t *testing.T) {
 	t.Run("found", func(t *testing.T) {
 		tests := map[string]struct {
-			key             string
-			caseInsensitive bool
-			structTags      []string
-			isInlineFuncs   []func(reflect.StructField) bool
-			v               interface{}
-			expect          interface{}
+			key                string
+			caseInsensitive    bool
+			structTags         []string
+			customExtractFuncs []func(ExtractFunc) ExtractFunc
+			isInlineFuncs      []func(reflect.StructField) bool
+			v                  interface{}
+			expect             interface{}
 		}{
 			"map[string]string": {
 				key: "key",
@@ -148,6 +149,34 @@ func TestKey_Extract(t *testing.T) {
 				},
 				expect: "xxx",
 			},
+			"struct (inline with custom extract funcs)": {
+				key:        "aaa",
+				structTags: []string{"json", "yaml"},
+				customExtractFuncs: []func(ExtractFunc) ExtractFunc{
+					func(f ExtractFunc) ExtractFunc {
+						return func(v reflect.Value) (reflect.Value, bool) {
+							if v.CanInterface() {
+								if vv, ok := v.Interface().(map[string]string); ok {
+									mp := map[string]any{}
+									for k, v := range vv {
+										mp[k] = v + v
+									}
+									if v, ok := f(reflect.ValueOf(&keyExtractorContext{v: mp})); ok {
+										return v, true
+									}
+								}
+							}
+							return f(v)
+						}
+					},
+				},
+				v: testTags{
+					M: map[string]string{
+						"aaa": "xxx",
+					},
+				},
+				expect: "xxxxxx",
+			},
 			"struct (custom inline func)": {
 				key: "aaa",
 				isInlineFuncs: []func(reflect.StructField) bool{
@@ -198,10 +227,11 @@ func TestKey_Extract(t *testing.T) {
 			test := test
 			t.Run(name, func(t *testing.T) {
 				e := &Key{
-					key:             test.key,
-					caseInsensitive: test.caseInsensitive,
-					structTags:      test.structTags,
-					isInlineFuncs:   test.isInlineFuncs,
+					key:                test.key,
+					caseInsensitive:    test.caseInsensitive,
+					structTags:         test.structTags,
+					customExtractFuncs: test.customExtractFuncs,
+					isInlineFuncs:      test.isInlineFuncs,
 				}
 				v, ok := e.Extract(reflect.ValueOf(test.v))
 				if !ok {

diff --git a/query.go b/query.go
@@ -45,11 +45,12 @@ func (q Query) Root() *Query {
 // Key is shorthand method to create Key and appends it.
 func (q Query) Key(k string) *Query {
 	return q.Append(&Key{
-		key:             k,
-		caseInsensitive: q.caseInsensitive,
-		structTags:      q.structTags,
-		fieldNameGetter: q.customStructFieldNameGetter,
-		isInlineFuncs:   q.customIsInlineFuncs,
+		key:                k,
+		caseInsensitive:    q.caseInsensitive,
+		structTags:         q.structTags,
+		customExtractFuncs: q.customExtractFuncs,
+		fieldNameGetter:    q.customStructFieldNameGetter,
+		isInlineFuncs:      q.customIsInlineFuncs,
 	})
 }