Fix #79, #80 -- Fix storing non-ASCII values on Python 2 and binary values on Python 3

3 files changed, 44 insertions, 9 deletions

diff --git a/ChangeLog b/ChangeLog
index 168685e..138e8ea 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -13,6 +13,9 @@
    *  Add flake8 testing and cleanups (PR from Tim Graham, cleanups from Sean
       Reifschneider) #112
 
+   * Fixed storing non-ASCII values on Python 2 and binary values on Python 3
+     (PR from Nicolas Noé) #135
+
 Fri, 27 May 2016 13:44:55 -0600  Sean Reifschneider  <jafo@tummy.com>
 
    *  1.58 release.
diff --git a/memcache.py b/memcache.py
index cf2a89a..efb8bff 100644
--- a/memcache.py
+++ b/memcache.py
@@ -134,6 +134,7 @@ class Client(threading.local):
     _FLAG_INTEGER = 1 << 1
     _FLAG_LONG = 1 << 2
     _FLAG_COMPRESSED = 1 << 3
+    _FLAG_TEXT = 1 << 4
 
     _SERVER_RETRIES = 10  # how many times to try finding a free server.
 
@@ -955,11 +956,16 @@ class Client(threading.local):
         the new value itself.
         """
         flags = 0
-        if isinstance(val, six.binary_type):
+        # Check against the exact type, rather than using isinstance(), so that
+        # subclasses of native types (such as markup-safe strings) are pickled
+        # and restored as instances of the correct class.
+        val_type = type(val)
+        if val_type == six.binary_type:
             pass
-        elif isinstance(val, six.text_type):
+        elif val_type == six.text_type:
+            flags |= Client._FLAG_TEXT
             val = val.encode('utf-8')
-        elif isinstance(val, int):
+        elif val_type == int:
             flags |= Client._FLAG_INTEGER
             val = '%d' % val
             if six.PY3:
@@ -1250,13 +1256,11 @@ class Client(threading.local):
         if flags & Client._FLAG_COMPRESSED:
             buf = self.decompressor(buf)
             flags &= ~Client._FLAG_COMPRESSED
-
         if flags == 0:
-            # Bare string
-            if six.PY3:
-                val = buf.decode('utf8')
-            else:
-                val = buf
+            # Bare bytes
+            val = buf
+        elif flags & Client._FLAG_TEXT:
+            val = buf.decode('utf-8')
         elif flags & Client._FLAG_INTEGER:
             val = int(buf)
         elif flags & Client._FLAG_LONG:
diff --git a/tests/test_memcache.py b/tests/test_memcache.py
index 58024cd..8fe0f66 100644
--- a/tests/test_memcache.py
+++ b/tests/test_memcache.py
@@ -1,6 +1,8 @@
+# -*- coding: utf-8 -*-
 from __future__ import print_function
 
 import unittest
+import zlib
 
 import mock
 import six
@@ -128,6 +130,32 @@ class TestMemcache(unittest.TestCase):
         value = self.mc.get(key)
         self.assertEqual(value, 5)
 
+    def test_unicode_value(self):
+        key = 'key'
+        value = u'Iñtërnâtiônàlizætiøn2'
+        self.mc.set(key, value)
+        cached_value = self.mc.get(key)
+        self.assertEqual(value, cached_value)
+
+    def test_binary_string(self):
+        value = 'value_to_be_compressed'
+        compressed_value = zlib.compress(value.encode())
+
+        self.mc.set('binary1', compressed_value)
+        compressed_result = self.mc.get('binary1')
+        self.assertEqual(compressed_value, compressed_result)
+        self.assertEqual(value, zlib.decompress(compressed_result).decode())
+
+        self.mc.add('binary1-add', compressed_value)
+        compressed_result = self.mc.get('binary1-add')
+        self.assertEqual(compressed_value, compressed_result)
+        self.assertEqual(value, zlib.decompress(compressed_result).decode())
+
+        self.mc.set_multi({'binary1-set_many': compressed_value})
+        compressed_result = self.mc.get('binary1-set_many')
+        self.assertEqual(compressed_value, compressed_result)
+        self.assertEqual(value, zlib.decompress(compressed_result).decode())
+
     def test_ignore_too_large_value(self):
         # NOTE: "MemCached: while expecting[...]" is normal...
         key = 'keyhere'