Skip to content

Commit

Permalink
Add ZMScore cmd
Browse files Browse the repository at this point in the history
  • Loading branch information
Benjamin Prieur committed Mar 8, 2021
1 parent f3a31a3 commit abb58ff
Show file tree
Hide file tree
Showing 3 changed files with 89 additions and 0 deletions.
49 changes: 49 additions & 0 deletions command.go
Original file line number Diff line number Diff line change
Expand Up @@ -810,6 +810,55 @@ func (cmd *FloatCmd) readReply(rd *proto.Reader) (err error) {

//------------------------------------------------------------------------------

type FloatSliceCmd struct {
baseCmd

val []float64
}

var _ Cmder = (*FloatSliceCmd)(nil)

func NewFloatSliceCmd(ctx context.Context, args ...interface{}) *FloatSliceCmd {
return &FloatSliceCmd{
baseCmd: baseCmd{
ctx: ctx,
args: args,
},
}
}

func (cmd *FloatSliceCmd) Val() []float64 {
return cmd.val
}

func (cmd *FloatSliceCmd) Result() ([]float64, error) {
return cmd.val, cmd.err
}

func (cmd *FloatSliceCmd) String() string {
return cmdString(cmd, cmd.val)
}

func (cmd *FloatSliceCmd) readReply(rd *proto.Reader) error {
_, err := rd.ReadArrayReply(func(rd *proto.Reader, n int64) (interface{}, error) {
cmd.val = make([]float64, n)
for i := 0; i < len(cmd.val); i++ {
switch num, err := rd.ReadFloatReply(); {
case err == Nil:
cmd.val[i] = 0
case err != nil:
return nil, err
default:
cmd.val[i] = num
}
}
return nil, nil
})
return err
}

//------------------------------------------------------------------------------

type StringSliceCmd struct {
baseCmd

Expand Down
13 changes: 13 additions & 0 deletions commands.go
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,7 @@ type Cmdable interface {
ZLexCount(ctx context.Context, key, min, max string) *IntCmd
ZIncrBy(ctx context.Context, key string, increment float64, member string) *FloatCmd
ZInterStore(ctx context.Context, destination string, store *ZStore) *IntCmd
ZMScore(ctx context.Context, key string, members ...string) *FloatSliceCmd
ZPopMax(ctx context.Context, key string, count ...int64) *ZSliceCmd
ZPopMin(ctx context.Context, key string, count ...int64) *ZSliceCmd
ZRange(ctx context.Context, key string, start, stop int64) *StringSliceCmd
Expand Down Expand Up @@ -2010,6 +2011,18 @@ func (c cmdable) ZInterStore(ctx context.Context, destination string, store *ZSt
return cmd
}

func (c cmdable) ZMScore(ctx context.Context, key string, members ...string) *FloatSliceCmd {
args := make([]interface{}, 2+len(members))
args[0] = "zmscore"
args[1] = key
for i, member := range members {
args[2+i] = member
}
cmd := NewFloatSliceCmd(ctx, args...)
_ = c(ctx, cmd)
return cmd
}

func (c cmdable) ZPopMax(ctx context.Context, key string, count ...int64) *ZSliceCmd {
args := []interface{}{
"zpopmax",
Expand Down
27 changes: 27 additions & 0 deletions commands_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3151,6 +3151,33 @@ var _ = Describe("Commands", func() {
}}))
})

It("should ZMScore", func() {
zmScore := client.ZMScore(ctx, "zset", "one", "three")
Expect(zmScore.Err()).NotTo(HaveOccurred())
Expect(zmScore.Val()).To(HaveLen(2))
Expect(zmScore.Val()[0]).To(Equal(float64(0)))

err := client.ZAdd(ctx, "zset", &redis.Z{Score: 1, Member: "one"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", &redis.Z{Score: 2, Member: "two"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", &redis.Z{Score: 3, Member: "three"}).Err()
Expect(err).NotTo(HaveOccurred())

zmScore = client.ZMScore(ctx, "zset", "one", "three")
Expect(zmScore.Err()).NotTo(HaveOccurred())
Expect(zmScore.Val()).To(HaveLen(2))
Expect(zmScore.Val()[0]).To(Equal(float64(1)))

zmScore = client.ZMScore(ctx, "zset", "four")
Expect(zmScore.Err()).NotTo(HaveOccurred())
Expect(zmScore.Val()).To(HaveLen(1))

zmScore = client.ZMScore(ctx, "zset", "four", "one")
Expect(zmScore.Err()).NotTo(HaveOccurred())
Expect(zmScore.Val()).To(HaveLen(2))
})

It("should ZPopMax", func() {
err := client.ZAdd(ctx, "zset", &redis.Z{
Score: 1,
Expand Down

0 comments on commit abb58ff

Please sign in to comment.