From abb58ff31f2dbd6c1a3a13ea6e73e3abc3f8f5e4 Mon Sep 17 00:00:00 2001
From: Benjamin Prieur <benjaminp@voodoo.com>
Date: Fri, 5 Mar 2021 19:02:53 +0100
Subject: [PATCH] Add ZMScore cmd

---
 command.go       | 49 ++++++++++++++++++++++++++++++++++++++++++++++++
 commands.go      | 13 +++++++++++++
 commands_test.go | 27 ++++++++++++++++++++++++++
 3 files changed, 89 insertions(+)

diff --git a/command.go b/command.go
index 6ca06344..b45a3edc 100644
--- a/command.go
+++ b/command.go
@@ -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
 
diff --git a/commands.go b/commands.go
index 2acbd837..61e306b7 100644
--- a/commands.go
+++ b/commands.go
@@ -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
@@ -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",
diff --git a/commands_test.go b/commands_test.go
index b39bc407..457fa057 100644
--- a/commands_test.go
+++ b/commands_test.go
@@ -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,