Add optional support for unicode keys

memcached's ASCII protocol supports unicode keys, so lets support them
as well. Since using unicode keys for memcache is uncommon and to
preserve the previous behavior disable support by default.

4 files changed, 94 insertions, 31 deletions

diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index ab72652..ae18c28 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -1,4 +1,5 @@
 # Copyright 2012 Pinterest.com
+# -*- coding: utf-8 -*-
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -92,6 +93,18 @@ class ClientTestMixin(object):
         with pytest.raises(MemcacheIllegalInputError):
             _set()
 
+    def test_set_unicode_key_ok(self):
+        client = self.make_client([b'STORED\r\n'], allow_unicode_keys=True)
+
+        result = client.set(u'\u0FFF', b'value', noreply=False)
+        assert result is True
+
+    def test_set_unicode_key_ok_snowman(self):
+        client = self.make_client([b'STORED\r\n'], allow_unicode_keys=True)
+
+        result = client.set('my☃', b'value', noreply=False)
+        assert result is True
+
     def test_set_unicode_char_in_middle_of_key(self):
         client = self.make_client([b'STORED\r\n'])
 
@@ -101,6 +114,15 @@ class ClientTestMixin(object):
         with pytest.raises(MemcacheIllegalInputError):
             _set()
 
+    def test_set_unicode_char_in_middle_of_key_snowman(self):
+        client = self.make_client([b'STORED\r\n'])
+
+        def _set():
+            client.set('my☃', b'value', noreply=False)
+
+        with pytest.raises(MemcacheIllegalInputError):
+            _set()
+
     def test_set_unicode_value(self):
         client = self.make_client([b''])
 
@@ -110,6 +132,12 @@ class ClientTestMixin(object):
         with pytest.raises(MemcacheIllegalInputError):
             _set()
 
+    def test_set_unicode_char_in_middle_of_key_ok(self):
+        client = self.make_client([b'STORED\r\n'], allow_unicode_keys=True)
+
+        result = client.set('helloworld_\xb1901520_%c3', b'value', noreply=False)
+        assert result is True
+
     def test_set_noreply(self):
         client = self.make_client([])
         result = client.set(b'key', b'value', noreply=True)
@@ -626,6 +654,15 @@ class TestClient(ClientTestMixin, unittest.TestCase):
         with pytest.raises(MemcacheClientError):
             client.get(b'x' * 251)
 
+    def test_too_long_unicode_key(self):
+        client = self.make_client([b'STORED\r\n'], allow_unicode_keys=True)
+
+        with pytest.raises(MemcacheClientError):
+            client.get('my☃'*150)
+
+        with pytest.raises(MemcacheClientError):
+            client.get(u'\u0FFF'*150)
+
     def test_key_contains_spae(self):
         client = self.make_client([b'END\r\n'])
         with pytest.raises(MemcacheClientError):
diff --git a/pymemcache/test/test_client_hash.py b/pymemcache/test/test_client_hash.py
index 393a191..f523f4a 100644
--- a/pymemcache/test/test_client_hash.py
+++ b/pymemcache/test/test_client_hash.py
@@ -12,8 +12,8 @@ import socket
 
 class TestHashClient(ClientTestMixin, unittest.TestCase):
 
-    def make_client_pool(self, hostname, mock_socket_values, serializer=None):
-        mock_client = Client(hostname, serializer=serializer)
+    def make_client_pool(self, hostname, mock_socket_values, serializer=None, **kwargs):
+        mock_client = Client(hostname, serializer=serializer, **kwargs)
         mock_client.sock = MockSocket(mock_socket_values)
         client = PooledClient(hostname, serializer=serializer)
         client.client_pool = pool.ObjectPool(lambda: mock_client)
@@ -28,7 +28,8 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
             s = '%s:%s' % (ip, current_port)
             c = self.make_client_pool(
                 (ip, current_port),
-                vals
+                vals,
+                **kwargs
             )
             client.clients[s] = c
             client.hasher.add_node(s)
diff --git a/pymemcache/test/test_integration.py b/pymemcache/test/test_integration.py
index 17b26fd..32b6b6f 100644
--- a/pymemcache/test/test_integration.py
+++ b/pymemcache/test/test_integration.py
@@ -1,4 +1,5 @@
 # Copyright 2012 Pinterest.com
+# -*- coding: utf-8 -*-
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,27 +24,47 @@ from pymemcache.exceptions import (
 )
 
 
+def get_set_helper(client, key, value, key2, value2):
+    result = client.get(key)
+    assert result is None
+
+    client.set(key, value, noreply=False)
+    result = client.get(key)
+    assert result == value
+
+    client.set(key2, value2, noreply=True)
+    result = client.get(key2)
+    assert result == value2
+
+    result = client.get_many([key, key2])
+    assert result == {key: value, key2: value2}
+
+    result = client.get_many([])
+    assert result == {}
+
+
 @pytest.mark.integration()
 def test_get_set(client_class, host, port, socket_module):
     client = client_class((host, port), socket_module=socket_module)
     client.flush_all()
 
-    result = client.get('key')
-    assert result is None
+    key = b'key'
+    value = b'value'
+    key2 = b'key2'
+    value2 = b'value2'
+    get_set_helper(client, key, value, key2, value2)
 
-    client.set(b'key', b'value', noreply=False)
-    result = client.get(b'key')
-    assert result == b'value'
-
-    client.set(b'key2', b'value2', noreply=True)
-    result = client.get(b'key2')
-    assert result == b'value2'
 
-    result = client.get_many([b'key', b'key2'])
-    assert result == {b'key': b'value', b'key2': b'value2'}
+@pytest.mark.integration()
+def test_get_set_unicode_key(client_class, host, port, socket_module):
+    client = client_class((host, port), socket_module=socket_module, allow_unicode_keys=True)
+    client.flush_all()
 
-    result = client.get_many([])
-    assert result == {}
+    key = u"こんにちは"
+    value = b'hello'
+    key2 = 'my☃'
+    value2 = b'value2'
+    get_set_helper(client, key, value, key2, value2)
 
 
 @pytest.mark.integration()
diff --git a/pymemcache/test/utils.py b/pymemcache/test/utils.py
index 4414031..a4a539e 100644
--- a/pymemcache/test/utils.py
+++ b/pymemcache/test/utils.py
@@ -26,12 +26,14 @@ class MockMemcacheClient(object):
                  timeout=None,
                  no_delay=False,
                  ignore_exc=False,
-                 default_noreply=True):
+                 default_noreply=True,
+                 allow_unicode_keys=False):
 
         self._contents = {}
 
         self.serializer = serializer
         self.deserializer = deserializer
+        self.allow_unicode_keys = allow_unicode_keys
 
         # Unused, but present for interface compatibility
         self.server = server
@@ -41,13 +43,14 @@ class MockMemcacheClient(object):
         self.ignore_exc = ignore_exc
 
     def get(self, key, default=None):
-        if isinstance(key, six.text_type):
-            raise MemcacheIllegalInputError(key)
-        if isinstance(key, six.string_types):
-            try:
-                key = key.encode('ascii')
-            except (UnicodeEncodeError, UnicodeDecodeError):
-                raise MemcacheIllegalInputError
+        if not self.allow_unicode_keys:
+            if isinstance(key, six.text_type):
+                raise MemcacheIllegalInputError(key)
+            if isinstance(key, six.string_types):
+                try:
+                    key = key.encode('ascii')
+                except (UnicodeEncodeError, UnicodeDecodeError):
+                    raise MemcacheIllegalInputError
 
         if key not in self._contents:
             return default
@@ -72,15 +75,16 @@ class MockMemcacheClient(object):
     get_multi = get_many
 
     def set(self, key, value, expire=0, noreply=True):
-        if isinstance(key, six.text_type):
-            raise MemcacheIllegalInputError(key)
+        if not self.allow_unicode_keys:
+            if isinstance(key, six.text_type):
+                raise MemcacheIllegalInputError(key)
+            if isinstance(key, six.string_types):
+                try:
+                    key = key.encode('ascii')
+                except (UnicodeEncodeError, UnicodeDecodeError):
+                    raise MemcacheIllegalInputError
         if isinstance(value, six.text_type):
             raise MemcacheIllegalInputError(value)
-        if isinstance(key, six.string_types):
-            try:
-                key = key.encode('ascii')
-            except (UnicodeEncodeError, UnicodeDecodeError):
-                raise MemcacheIllegalInputError
         if isinstance(value, six.string_types):
             try:
                 value = value.encode('ascii')