allow the user to choose a connection class via the ConnectionPool

3 files changed, 33 insertions, 22 deletions

diff --git a/CHANGES b/CHANGES
index acf33e6..97ab91f 100644
--- a/CHANGES
+++ b/CHANGES
@@ -2,6 +2,10 @@
     * Added support for Hiredis. To use, simply "pip install hiredis" or 
       "easy_install hiredis". Thanks for Pieter Noordhuis for the hiredis-py
       bindings and the patch to redis-py.
+    * The connection class is chosed based on whether hiredis is installed
+      or not. To force the use of the PythonConnection, simply create
+      your own ConnectionPool instance with the connection_class argument
+      assigned to to PythonConnection class.
 * 2.2.2
     * Fixed a bug in ZREVRANK where retriving the rank of a value not in
       the zset would raise an error.
diff --git a/redis/connection.py b/redis/connection.py
index 3f65ccd..648f7a1 100644
--- a/redis/connection.py
+++ b/redis/connection.py
@@ -3,28 +3,6 @@ import socket
 import threading
 from redis.exceptions import ConnectionError, ResponseError, InvalidResponse
 
-class ConnectionPool(threading.local):
-    "Manages a list of connections on the local thread"
-    def __init__(self):
-        self.connections = {}
-
-    def make_connection_key(self, host, port, db):
-        "Create a unique key for the specified host, port and db"
-        return '%s:%s:%s' % (host, port, db)
-
-    def get_connection(self, host, port, db, password, socket_timeout):
-        "Return a specific connection for the specified host, port and db"
-        key = self.make_connection_key(host, port, db)
-        if key not in self.connections:
-            self.connections[key] = Connection(
-                host, port, db, password, socket_timeout)
-        return self.connections[key]
-
-    def get_all_connections(self):
-        "Return a list of all connection objects the manager knows about"
-        return self.connections.values()
-
-
 class BaseConnection(object):
     "Manages TCP communication to and from a Redis server"
     def __init__(self, host='localhost', port=6379, db=0, password=None,
@@ -190,3 +168,25 @@ try:
     Connection = HiredisConnection
 except ImportError:
     Connection = PythonConnection
+
+class ConnectionPool(threading.local):
+    "Manages a list of connections on the local thread"
+    def __init__(self, connection_class=None):
+        self.connections = {}
+        self.connection_class = connection_class or Connection
+
+    def make_connection_key(self, host, port, db):
+        "Create a unique key for the specified host, port and db"
+        return '%s:%s:%s' % (host, port, db)
+
+    def get_connection(self, host, port, db, password, socket_timeout):
+        "Return a specific connection for the specified host, port and db"
+        key = self.make_connection_key(host, port, db)
+        if key not in self.connections:
+            self.connections[key] = self.connection_class(
+                host, port, db, password, socket_timeout)
+        return self.connections[key]
+
+    def get_all_connections(self):
+        "Return a list of all connection objects the manager knows about"
+        return self.connections.values()
diff --git a/tests/__init__.py b/tests/__init__.py
index 8931b07..ab94711 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -4,6 +4,13 @@ from connection_pool import ConnectionPoolTestCase
 from pipeline import PipelineTestCase
 from lock import LockTestCase
 
+use_hiredis = False
+try:
+    import hiredis
+    use_hiredis = True
+except ImportError:
+    pass
+
 def all_tests():
     suite = unittest.TestSuite()
     suite.addTest(unittest.makeSuite(ServerCommandsTestCase))