Refactor set loop type: closes #84

rethinkdb/__init__.py

@@ -11,8 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
+
+import imp
 
 from rethinkdb import errors, version
+from rethinkdb import net
+import pkg_resources
 
 
 # The builtins here defends against re-importing something obscuring `object`.
@@ -43,3 +48,31 @@ def __init__(self):
         for module in (net, query, ast, errors):
             for function_name in module.__all__:
                 setattr(self, function_name, getattr(module, function_name))
+
+        self.set_loop_type(None)
+
+    def set_loop_type(self, library=None):
+        if library is None:
+            self.connection_type = net.DefaultConnection
+            return
+
+        # find module file
+        manager = pkg_resources.ResourceManager()
+        libPath = '%(library)s_net/net_%(library)s.py' % {'library': library}
+        if not manager.resource_exists(__name__, libPath):
+            raise ValueError('Unknown loop type: %r' % library)
+
+        # load the module
+        modulePath = manager.resource_filename(__name__, libPath)
+        moduleName = 'net_%s' % library
+        moduleFile, pathName, desc = imp.find_module(moduleName, [os.path.dirname(modulePath)])
+        module = imp.load_module('rethinkdb.' + moduleName, moduleFile, pathName, desc)
+
+        # set the connection type
+        self.connection_type = module.Connection
+
+        # cleanup
+        manager.cleanup_resources()
+
+    def connect(self, *args, **kwargs):
+        return self.make_connection(self.connection_type, *args, **kwargs)

rethinkdb/net.py

@@ -18,9 +18,7 @@
 
 import collections
 import errno
-import imp
 import numbers
-import os
 import pprint
 import socket
 import ssl
@@ -48,7 +46,7 @@
 from rethinkdb.handshake import HandshakeV1_0
 from rethinkdb.logger import default_logger
 
-__all__ = ['connect', 'set_loop_type', 'Connection', 'Cursor', 'DEFAULT_PORT']
+__all__ = ['Connection', 'Cursor', 'DEFAULT_PORT', 'DefaultConnection', 'make_connection']
 
 
 DEFAULT_PORT = 28015
@@ -705,10 +703,11 @@ def __init__(self, *args, **kwargs):
         Connection.__init__(self, ConnectionInstance, *args, **kwargs)
 
 
-connection_type = DefaultConnection
 
 
-def connect(
+
+def make_connection(
+        connection_type,
         host=None,
         port=None,
         db=None,
@@ -734,26 +733,3 @@ def connect(
 
     conn = connection_type(host, port, db, auth_key, user, password, timeout, ssl, _handshake_version, **kwargs)
     return conn.reconnect(timeout=timeout)
-
-
-def set_loop_type(library):
-    global connection_type
-    import pkg_resources
-
-    # find module file
-    manager = pkg_resources.ResourceManager()
-    libPath = '%(library)s_net/net_%(library)s.py' % {'library': library}
-    if not manager.resource_exists(__name__, libPath):
-        raise ValueError('Unknown loop type: %r' % library)
-
-    # load the module
-    modulePath = manager.resource_filename(__name__, libPath)
-    moduleName = 'net_%s' % library
-    moduleFile, pathName, desc = imp.find_module(moduleName, [os.path.dirname(modulePath)])
-    module = imp.load_module('rethinkdb.' + moduleName, moduleFile, pathName, desc)
-
-    # set the connection type
-    connection_type = module.Connection
-
-    # cleanup
-    manager.cleanup_resources()

rethinkdb/utils_common.py

@@ -67,7 +67,7 @@ def conn(self, test_connection=True):
 
         # cache a new connection
         if not os.getpid() in self.__local.connCache:
-            self.__local.connCache[os.getpid()] = net.connect(**self.__connectOptions)
+            self.__local.connCache[os.getpid()] = net.make_connection(net.DefaultConnection, **self.__connectOptions)
 
         # return the connection
         return self.__local.connCache[os.getpid()]