From 519f98e3704c03505283d18c28720ea9d520fdad Mon Sep 17 00:00:00 2001
From: Joe Gordon <jogo@pinterest.com>
Date: Thu, 6 Sep 2018 17:08:04 -0700
Subject: [PATCH] Fix support newbytes from future

Previously python2 code using python-future to backport the py3 bytes
behavior would trigger the following exception:

code:

  from builtins import bytes as newbytes
  from pymemcache.client.base import Client
  client = Client(('localhost', 11211))
  client.set(newbytes('key'), 'value')

  Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
    File "/usr/local/lib/python2.7/site-packages/pymemcache/client/base.py", line 297, in set
      return self._store_cmd(b'set', key, expire, noreply, value)
    File "/usr/local/lib/python2.7/site-packages/pymemcache/client/base.py", line 770, in _store_cmd
      key = self.check_key(key)
    File "/usr/local/lib/python2.7/site-packages/pymemcache/client/base.py", line 251, in check_key
      key_prefix=self.key_prefix)
    File "/usr/local/lib/python2.7/site-packages/pymemcache/client/base.py", line 91, in _check_key
      key = key.encode('ascii')
    File "/usr/local/lib/python2.7/site-packages/future/types/newbytes.py", line 381, in __getattribute__
      raise AttributeError("encode method has been disabled in newbytes")
  AttributeError: encode method has been disabled in newbytes

Add a test case for this and fix.
---
 pymemcache/client/base.py           |  5 ++++-
 pymemcache/test/test_client.py      |  7 +++++++
 pymemcache/test/test_integration.py |  5 +++++
 pymemcache/test/utils.py            | 10 ++++++++--
 test-requirements.txt               |  1 +
 5 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 6e9de0fc..846addae 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -88,7 +88,10 @@ def _check_key(key, allow_unicode_keys, key_prefix=b''):
             key = key.encode('utf8')
     elif isinstance(key, VALID_STRING_TYPES):
         try:
-            key = key.encode('ascii')
+            if isinstance(key, bytes):
+                key = key.decode().encode('ascii')
+            else:
+                key = key.encode('ascii')
         except (UnicodeEncodeError, UnicodeDecodeError):
             raise MemcacheIllegalInputError("Non-ASCII key: '%r'" % key)
 
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index e7871367..eaadadd3 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from builtins import bytes as newbytes
+
 import collections
 import errno
 import functools
@@ -100,6 +102,11 @@ def test_set_success(self):
         result = client.set(b'key', b'value', noreply=False)
         assert result is True
 
+    def test_set_future(self):
+        client = self.make_client([b'STORED\r\n'])
+        result = client.set(newbytes(b'key'), newbytes(b'value'), noreply=False)
+        assert result is True
+
     def test_set_unicode_key(self):
         client = self.make_client([b''])
 
diff --git a/pymemcache/test/test_integration.py b/pymemcache/test/test_integration.py
index 4cbf30eb..a5d295f9 100644
--- a/pymemcache/test/test_integration.py
+++ b/pymemcache/test/test_integration.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from builtins import bytes as newbytes
+
 from collections import defaultdict
 import json
 import pytest
@@ -41,6 +43,9 @@ def get_set_helper(client, key, value, key2, value2):
     result = client.get(key2)
     assert result == value2
 
+    result = client.get(newbytes(key2))
+    assert result == value2
+
     result = client.get_many([key, key2])
     assert result == {key: value, key2: value2}
 
diff --git a/pymemcache/test/utils.py b/pymemcache/test/utils.py
index 7325f2d9..0584772f 100644
--- a/pymemcache/test/utils.py
+++ b/pymemcache/test/utils.py
@@ -48,7 +48,10 @@ def get(self, key, default=None):
                 raise MemcacheIllegalInputError(key)
             if isinstance(key, six.string_types):
                 try:
-                    key = key.encode('ascii')
+                    if isinstance(key, bytes):
+                        key = key.decode().encode('ascii')
+                    else:
+                        key = key.encode('ascii')
                 except (UnicodeEncodeError, UnicodeDecodeError):
                     raise MemcacheIllegalInputError
 
@@ -80,7 +83,10 @@ def set(self, key, value, expire=0, noreply=True):
                 raise MemcacheIllegalInputError(key)
             if isinstance(key, six.string_types):
                 try:
-                    key = key.encode('ascii')
+                    if isinstance(key, bytes):
+                        key = key.decode().encode('ascii')
+                    else:
+                        key = key.encode('ascii')
                 except (UnicodeEncodeError, UnicodeDecodeError):
                     raise MemcacheIllegalInputError
         if isinstance(value, six.text_type):
diff --git a/test-requirements.txt b/test-requirements.txt
index 7bcdf2e6..5c3424ad 100644
--- a/test-requirements.txt
+++ b/test-requirements.txt
@@ -4,3 +4,4 @@ pytest-cov
 gevent==1.3.6; "PyPy" not in platform_python_implementation
 pylibmc
 python-memcached
+future