feat: field mask util.

This commit is contained in:
tx7do
2023-11-06 20:15:07 +08:00
parent 413e14ac78
commit 95ce578ddf
5 changed files with 278 additions and 1 deletions

View File

@@ -0,0 +1,210 @@
package fieldmaskutil
import (
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/reflect/protoreflect"
)
// Filter keeps the msg fields that are listed in the paths and clears all the rest.
//
// This is a handy wrapper for NestedMask.Filter method.
// If the same paths are used to process multiple proto messages use NestedMask.Filter method directly.
func Filter(msg proto.Message, paths []string) {
NestedMaskFromPaths(paths).Filter(msg)
}
// Prune clears all the fields listed in paths from the given msg.
//
// This is a handy wrapper for NestedMask.Prune method.
// If the same paths are used to process multiple proto messages use NestedMask.Filter method directly.
func Prune(msg proto.Message, paths []string) {
NestedMaskFromPaths(paths).Prune(msg)
}
// Overwrite overwrites all the fields listed in paths in the dest msg using values from src msg.
//
// This is a handy wrapper for NestedMask.Overwrite method.
// If the same paths are used to process multiple proto messages use NestedMask.Overwrite method directly.
func Overwrite(src, dest proto.Message, paths []string) {
NestedMaskFromPaths(paths).Overwrite(src, dest)
}
// NestedMask represents a field mask as a recursive map.
type NestedMask map[string]NestedMask
// NestedMaskFromPaths creates an instance of NestedMask for the given paths.
func NestedMaskFromPaths(paths []string) NestedMask {
mask := make(NestedMask)
for _, path := range paths {
curr := mask
var letters []rune
for _, letter := range path {
if letter == '.' {
if len(letters) == 0 {
continue
}
key := string(letters)
c, ok := curr[key]
if !ok {
c = make(NestedMask)
curr[key] = c
}
curr = c
letters = nil
continue
}
letters = append(letters, letter)
}
if len(letters) != 0 {
key := string(letters)
if _, ok := curr[key]; !ok {
curr[key] = make(NestedMask)
}
}
}
return mask
}
// Filter keeps the msg fields that are listed in the paths and clears all the rest.
//
// If the mask is empty then all the fields are kept.
// Paths are assumed to be valid and normalized otherwise the function may panic.
// See google.golang.org/protobuf/types/known/fieldmaskpb for details.
func (mask NestedMask) Filter(msg proto.Message) {
if len(mask) == 0 {
return
}
rft := msg.ProtoReflect()
rft.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
m, ok := mask[string(fd.Name())]
if ok {
if len(m) == 0 {
return true
}
if fd.IsMap() {
xmap := rft.Get(fd).Map()
xmap.Range(func(mk protoreflect.MapKey, mv protoreflect.Value) bool {
if mi, ok := m[mk.String()]; ok {
if i, ok := mv.Interface().(protoreflect.Message); ok && len(mi) > 0 {
mi.Filter(i.Interface())
}
} else {
xmap.Clear(mk)
}
return true
})
} else if fd.IsList() {
list := rft.Get(fd).List()
for i := 0; i < list.Len(); i++ {
m.Filter(list.Get(i).Message().Interface())
}
} else if fd.Kind() == protoreflect.MessageKind {
m.Filter(rft.Get(fd).Message().Interface())
}
} else {
rft.Clear(fd)
}
return true
})
}
// Prune clears all the fields listed in paths from the given msg.
//
// All other fields are kept untouched. If the mask is empty no fields are cleared.
// This operation is the opposite of NestedMask.Filter.
// Paths are assumed to be valid and normalized otherwise the function may panic.
// See google.golang.org/protobuf/types/known/fieldmaskpb for details.
func (mask NestedMask) Prune(msg proto.Message) {
if len(mask) == 0 {
return
}
rft := msg.ProtoReflect()
rft.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
m, ok := mask[string(fd.Name())]
if ok {
if len(m) == 0 {
rft.Clear(fd)
return true
}
if fd.IsMap() {
xmap := rft.Get(fd).Map()
xmap.Range(func(mk protoreflect.MapKey, mv protoreflect.Value) bool {
if mi, ok := m[mk.String()]; ok {
if i, ok := mv.Interface().(protoreflect.Message); ok && len(mi) > 0 {
mi.Prune(i.Interface())
} else {
xmap.Clear(mk)
}
}
return true
})
} else if fd.IsList() {
list := rft.Get(fd).List()
for i := 0; i < list.Len(); i++ {
m.Prune(list.Get(i).Message().Interface())
}
} else if fd.Kind() == protoreflect.MessageKind {
m.Prune(rft.Get(fd).Message().Interface())
}
}
return true
})
}
// Overwrite overwrites all the fields listed in paths in the dest msg using values from src msg.
//
// All other fields are kept untouched. If the mask is empty, no fields are overwritten.
// Supports scalars, messages, repeated fields, and maps.
// If the parent of the field is nil message, the parent is initiated before overwriting the field
// If the field in src is empty value, the field in dest is cleared.
// Paths are assumed to be valid and normalized otherwise the function may panic.
func (mask NestedMask) Overwrite(src, dest proto.Message) {
mask.overwrite(src.ProtoReflect(), dest.ProtoReflect())
}
func (mask NestedMask) overwrite(src, dest protoreflect.Message) {
for k, v := range mask {
srcFD := src.Descriptor().Fields().ByName(protoreflect.Name(k))
destFD := dest.Descriptor().Fields().ByName(protoreflect.Name(k))
if srcFD == nil || destFD == nil {
continue
}
// Leaf mask -> copy value from src to dest
if len(v) == 0 {
if srcFD.Kind() == destFD.Kind() { // TODO: Full type equality check
val := src.Get(srcFD)
if isValid(srcFD, val) {
dest.Set(destFD, val)
} else {
dest.Clear(destFD)
}
}
} else if srcFD.Kind() == protoreflect.MessageKind {
// If dest field is nil
if !dest.Get(destFD).Message().IsValid() {
dest.Set(destFD, protoreflect.ValueOf(dest.Get(destFD).Message().New()))
}
v.overwrite(src.Get(srcFD).Message(), dest.Get(destFD).Message())
}
}
}
func isValid(fd protoreflect.FieldDescriptor, val protoreflect.Value) bool {
if fd.IsMap() {
return val.Map().IsValid()
} else if fd.IsList() {
return val.List().IsValid()
} else if fd.Message() != nil {
return val.Message().IsValid()
}
return true
}

View File

@@ -0,0 +1,59 @@
package fieldmaskutil
import (
"reflect"
"testing"
)
func Test_NestedMaskFromPaths(t *testing.T) {
type args struct {
paths []string
}
tests := []struct {
name string
args args
want NestedMask
}{
{
name: "no nested fields",
args: args{paths: []string{"a", "b", "c"}},
want: NestedMask{"a": NestedMask{}, "b": NestedMask{}, "c": NestedMask{}},
},
{
name: "with nested fields",
args: args{paths: []string{"aaa.bb.c", "dd.e", "f"}},
want: NestedMask{
"aaa": NestedMask{"bb": NestedMask{"c": NestedMask{}}},
"dd": NestedMask{"e": NestedMask{}},
"f": NestedMask{}},
},
{
name: "single field",
args: args{paths: []string{"a"}},
want: NestedMask{"a": NestedMask{}},
},
{
name: "empty fields",
args: args{paths: []string{}},
want: NestedMask{},
},
{
name: "invalid input",
args: args{paths: []string{".", "..", "..."}},
want: NestedMask{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := NestedMaskFromPaths(tt.args.paths); !reflect.DeepEqual(got, tt.want) {
t.Errorf("NestedMaskFromPaths() = %v, want %v", got, tt.want)
}
})
}
}
func BenchmarkNestedMaskFromPaths(b *testing.B) {
for i := 0; i < b.N; i++ {
NestedMaskFromPaths([]string{"aaa.bbb.c.d.e.f", "aa.b.cc.ddddddd", "e", "f", "g.h.i.j.k"})
}
}

1
go.mod
View File

@@ -10,6 +10,7 @@ require (
github.com/stretchr/testify v1.8.4
golang.org/x/crypto v0.14.0
golang.org/x/exp v0.0.0-20231006140011-7918f672742d
google.golang.org/protobuf v1.31.0
)
require (

7
go.sum
View File

@@ -3,6 +3,9 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/gobwas/glob v0.2.3 h1:A4xDbljILXROh+kObIiy5kIaPYD8e96x1tgBhUI5J+Y=
github.com/gobwas/glob v0.2.3/go.mod h1:d3Ez4x06l9bZtSvzIay5+Yzi0fmZzPgnTbPcKjJAkT8=
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg=
github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4=
github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/gosimple/slug v1.13.1 h1:bQ+kpX9Qa6tHRaK+fZR0A0M2Kd7Pa5eHPPsb1JpHD+Q=
@@ -30,6 +33,10 @@ golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

View File

@@ -1,4 +1,4 @@
git tag v1.1.6
git tag v1.1.7
git tag bank_card/v1.1.0
git tag entgo/v1.1.8
git tag geoip/v1.1.0