Fix flags when setting multiple values at once (#248)v2.2.1

We introduced the ability to override the serializer-returned flags
values in 26f7c1b1. Unfortunately, there was a flaw in that logic which
resulted in the first item's flags being used for all later items. This
bug primarily affected set_many()'s behavior when the data dictionary
contained multiple different value types which were assigned different
per-value flags values by the serializer.

For example:

    set_many({'a': 'string', 'b': 10})

If a serializer returned different flags for strings (e.g. 1) and
integer values (e.g. 2), the previous logic would have set ``flags``
to 1 the first time through the loop and repeated that value for the
second item, instead of using 2.

This was the intended behavior when ``flags`` was explicitly passed to
set_many(), but not for the default case where we still want to respect
the flags values returned by the serializer.

4 files changed, 47 insertions, 10 deletions

diff --git a/ChangeLog.rst b/ChangeLog.rst
index ccb30ee..c8d6376 100644
--- a/ChangeLog.rst
+++ b/ChangeLog.rst
@@ -1,6 +1,10 @@
 Change Log
 ==========
 
+New in version 2.2.1
+--------------------
+* Fix ``flags`` when setting multiple differently-typed values at once.
+
 New in version 2.2.0
 --------------------
 * Drop official support for Python 3.4.
diff --git a/pymemcache/__init__.py b/pymemcache/__init__.py
index 984d042..ae051ba 100644
--- a/pymemcache/__init__.py
+++ b/pymemcache/__init__.py
@@ -1,4 +1,4 @@
-__version__ = '2.2.0'
+__version__ = '2.2.1'
 
 from pymemcache.client.base import Client  # noqa
 from pymemcache.client.base import PooledClient  # noqa
diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 8dd1191..02c70d9 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -854,15 +854,14 @@ class Client(object):
 
             key = self.check_key(key)
             if self.serializer:
-                data, serializer_flags = self.serializer(key, data)
-                # Use the serializer's flags if the caller hasn't specified an
-                # explicit, overridding value.
-                if flags is None:
-                    flags = serializer_flags
+                data, data_flags = self.serializer(key, data)
+            else:
+                data_flags = 0
 
-            # Set flags to 0 if none were provided by the caller or serializer.
-            if flags is None:
-                flags = 0
+            # If 'flags' was explicitly provided, it overrides the value
+            # returned by the serializer.
+            if flags is not None:
+                data_flags = flags
 
             if not isinstance(data, six.binary_type):
                 try:
@@ -872,7 +871,7 @@ class Client(object):
                             "Data values must be binary-safe: %s" % e)
 
             cmds.append(name + b' ' + key + b' ' +
-                        six.text_type(flags).encode(self.encoding) +
+                        six.text_type(data_flags).encode(self.encoding) +
                         b' ' + expire +
                         b' ' + six.text_type(len(data)).encode(self.encoding) +
                         extra + b'\r\n' + data + b'\r\n')
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index 773bac9..0cebf32 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -674,6 +674,40 @@ class TestClient(ClientTestMixin, unittest.TestCase):
             b'set key 0 0 10 noreply\r\n{"c": "d"}\r\n'
         ]
 
+    def test_serialization_flags(self):
+        def _ser(key, value):
+            return value, 1 if isinstance(value, int) else 0
+
+        client = self.make_client(
+            [b'STORED\r\n', b'STORED\r\n'], serializer=_ser)
+        client.set_many(
+            collections.OrderedDict([(b'a', b's'), (b'b', 0)]), noreply=False)
+        assert client.sock.send_bufs == [
+            b'set a 0 0 1\r\ns\r\nset b 1 0 1\r\n0\r\n'
+        ]
+
+    def test_serialization_overridden_flags(self):
+        def _ser(key, value):
+            return value, 1 if isinstance(value, int) else 0
+
+        client = self.make_client(
+            [b'STORED\r\n', b'STORED\r\n'], serializer=_ser)
+        client.set_many(
+            collections.OrderedDict([(b'a', b's'), (b'b', 0)]),
+            noreply=False, flags=5)
+        assert client.sock.send_bufs == [
+            b'set a 5 0 1\r\ns\r\nset b 5 0 1\r\n0\r\n'
+        ]
+
+    def test_explicit_flags(self):
+        client = self.make_client([b'STORED\r\n', b'STORED\r\n'])
+        client.set_many(
+            collections.OrderedDict([(b'a', b's'), (b'b', 0)]),
+            noreply=False, flags=5)
+        assert client.sock.send_bufs == [
+            b'set a 5 0 1\r\ns\r\nset b 5 0 1\r\n0\r\n'
+        ]
+
     def test_set_socket_handling(self):
         client = self.make_client([b'STORED\r\n'])
         result = client.set(b'key', b'value', noreply=False)