add an Encoder object responsible for encoding/decoding bytes and strings

this simplifies multiple places that needs to encode and decode values

1 files changed, 49 insertions, 39 deletions

diff --git a/redis/connection.py b/redis/connection.py
index 77eb44d..c8f0513 100755
--- a/redis/connection.py
+++ b/redis/connection.py
@@ -94,6 +94,38 @@ class Token(object):
         return self.value
 
 
+class Encoder(object):
+    "Encode strings to bytes and decode bytes to strings"
+
+    def __init__(self, encoding, encoding_errors, decode_responses):
+        self.encoding = encoding
+        self.encoding_errors = encoding_errors
+        self.decode_responses = decode_responses
+
+    def encode(self, value):
+        "Return a bytestring representation of the value"
+        if isinstance(value, Token):
+            return value.encoded_value
+        elif isinstance(value, bytes):
+            return value
+        elif isinstance(value, (int, long)):
+            value = b(str(value))
+        elif isinstance(value, float):
+            value = b(repr(value))
+        elif not isinstance(value, basestring):
+            # an object we don't know how to deal with. default to unicode()
+            value = unicode(value)
+        if isinstance(value, unicode):
+            value = value.encode(self.encoding, self.encoding_errors)
+        return value
+
+    def decode(self, value, force=False):
+        "Return a unicode string from the byte representation"
+        if (self.decode_responses or force) and isinstance(value, bytes):
+            value = value.decode(self.encoding, self.encoding_errors)
+        return value
+
+
 class BaseParser(object):
     EXCEPTION_CLASSES = {
         'ERR': {
@@ -217,10 +249,9 @@ class SocketBuffer(object):
 
 class PythonParser(BaseParser):
     "Plain Python parsing class"
-    encoding = None
-
     def __init__(self, socket_read_size):
         self.socket_read_size = socket_read_size
+        self.encoder = None
         self._sock = None
         self._buffer = None
 
@@ -234,8 +265,7 @@ class PythonParser(BaseParser):
         "Called when the socket connects"
         self._sock = connection._sock
         self._buffer = SocketBuffer(self._sock, self.socket_read_size)
-        if connection.decode_responses:
-            self.encoding = connection.encoding
+        self.encoder = connection.encoder
 
     def on_disconnect(self):
         "Called when the socket disconnects"
@@ -245,7 +275,7 @@ class PythonParser(BaseParser):
         if self._buffer is not None:
             self._buffer.close()
             self._buffer = None
-        self.encoding = None
+        self.encoder = None
 
     def can_read(self):
         return self._buffer and bool(self._buffer.length)
@@ -292,8 +322,8 @@ class PythonParser(BaseParser):
             if length == -1:
                 return None
             response = [self.read_response() for i in xrange(length)]
-        if isinstance(response, bytes) and self.encoding:
-            response = response.decode(self.encoding)
+        if isinstance(response, bytes):
+            response = self.encoder.decode(response)
         return response
 
 
@@ -324,8 +354,8 @@ class HiredisParser(BaseParser):
         if not HIREDIS_SUPPORTS_CALLABLE_ERRORS:
             kwargs['replyError'] = ResponseError
 
-        if connection.decode_responses:
-            kwargs['encoding'] = connection.encoding
+        if connection.encoder.decode_responses:
+            kwargs['encoding'] = connection.encoder.encoding
         self._reader = hiredis.Reader(**kwargs)
         self._next_response = False
 
@@ -421,9 +451,7 @@ class Connection(object):
         self.socket_keepalive = socket_keepalive
         self.socket_keepalive_options = socket_keepalive_options or {}
         self.retry_on_timeout = retry_on_timeout
-        self.encoding = encoding
-        self.encoding_errors = encoding_errors
-        self.decode_responses = decode_responses
+        self.encoder = Encoder(encoding, encoding_errors, decode_responses)
         self._sock = None
         self._parser = parser_class(socket_read_size=socket_read_size)
         self._description_args = {
@@ -601,22 +629,6 @@ class Connection(object):
             raise response
         return response
 
-    def encode(self, value):
-        "Return a bytestring representation of the value"
-        if isinstance(value, Token):
-            return value.encoded_value
-        elif isinstance(value, bytes):
-            return value
-        elif isinstance(value, (int, long)):
-            value = b(str(value))
-        elif isinstance(value, float):
-            value = b(repr(value))
-        elif not isinstance(value, basestring):
-            value = unicode(value)
-        if isinstance(value, unicode):
-            value = value.encode(self.encoding, self.encoding_errors)
-        return value
-
     def pack_command(self, *args):
         "Pack a series of arguments into the Redis protocol"
         output = []
@@ -635,7 +647,7 @@ class Connection(object):
         buff = SYM_EMPTY.join(
             (SYM_STAR, b(str(len(args))), SYM_CRLF))
 
-        for arg in imap(self.encode, args):
+        for arg in imap(self.encoder.encode, args):
             # to avoid large string mallocs, chunk the command into the
             # output list if we're sending large values
             if len(buff) > 6000 or len(arg) > 6000:
@@ -724,9 +736,7 @@ class UnixDomainSocketConnection(Connection):
         self.password = password
         self.socket_timeout = socket_timeout
         self.retry_on_timeout = retry_on_timeout
-        self.encoding = encoding
-        self.encoding_errors = encoding_errors
-        self.decode_responses = decode_responses
+        self.encoder = Encoder(encoding, encoding_errors, decode_responses)
         self._sock = None
         self._parser = parser_class(socket_read_size=socket_read_size)
         self._description_args = {
@@ -956,14 +966,14 @@ class ConnectionPool(object):
         self._in_use_connections.add(connection)
         return connection
 
-    def get_encoding(self):
-        "Return information about the encoding settings"
+    def get_encoder(self):
+        "Return an encoder based on encoding settings"
         kwargs = self.connection_kwargs
-        return {
-            'encoding': kwargs.get('encoding', 'utf-8'),
-            'encoding_errors': kwargs.get('encoding_errors', 'strict'),
-            'decode_responses': kwargs.get('decode_responses', False)
-        }
+        return Encoder(
+            encoding=kwargs.get('encoding', 'utf-8'),
+            encoding_errors=kwargs.get('encoding_errors', 'strict'),
+            decode_responses=kwargs.get('decode_responses', False)
+        )
 
     def make_connection(self):
         "Create a new connection"