add document_class parameter/property to Connection

Mike Dirolf · Mike Dirolf · commit afa57fa97f70 · 2010-05-18T13:41:28.000-04:00
diff --git a/doc/api/pymongo/connection.rst b/doc/api/pymongo/connection.rst
@@ -4,7 +4,7 @@
 .. automodule:: pymongo.connection
    :synopsis: Tools for connecting to MongoDB
 
-   .. autoclass:: pymongo.connection.Connection([host='localhost'[, port=27017[, pool_size=None[, auto_start_request=None[, timeout=None[, slave_okay=False[, network_timeout=None]]]]]]])
+   .. autoclass:: pymongo.connection.Connection([host='localhost'[, port=27017[, pool_size=None[, auto_start_request=None[, timeout=None[, slave_okay=False[, network_timeout=None[, document_class=dict]]]]]]]])
 
       .. automethod:: from_uri([uri='mongodb://localhost'])
       .. automethod:: paired(left[, right=('localhost', 27017)])
@@ -19,6 +19,7 @@
       .. autoattribute:: host
       .. autoattribute:: port
       .. autoattribute:: slave_okay
+      .. autoattribute:: document_class
       .. automethod:: database_names
       .. automethod:: drop_database
       .. automethod:: copy_database(from_name, to_name[, from_host=None[, username=None[, password=None]]])
diff --git a/doc/api/pymongo/cursor.rst b/doc/api/pymongo/cursor.rst
@@ -4,7 +4,7 @@
 .. automodule:: pymongo.cursor
    :synopsis: Tools for iterating over MongoDB query results
 
-   .. autoclass:: pymongo.cursor.Cursor(collection, spec, fields, skip, limit, slave_okay, timeout, tailable[, snapshot=False[, sort=None[, max_scan=None]]])
+   .. autoclass:: pymongo.cursor.Cursor
       :members:
 
       .. describe:: c[index]
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -460,7 +460,8 @@ def find(self, *args, **kwargs):
           - `max_scan` (optional): limit the number of documents
             examined when performing the query
           - `as_class` (optional): class to use for documents in the
-            query result (default is :class:`dict`)
+            query result (default is
+            :attr:`~pymongo.connection.Connection.document_class`)
 
         .. note:: The `max_scan` parameter requires server
            version **>= 1.5.1**
diff --git a/pymongo/connection.py b/pymongo/connection.py
@@ -108,7 +108,7 @@ class Connection(object): # TODO support auth for pooling
 
     def __init__(self, host=None, port=None, pool_size=None,
                  auto_start_request=None, timeout=None, slave_okay=False,
-                 network_timeout=None, _connect=True):
+                 network_timeout=None, document_class=dict, _connect=True):
         """Create a new connection to a single MongoDB instance at *host:port*.
 
         The resultant connection object has connection-pooling built in. It
@@ -131,13 +131,17 @@ def __init__(self, host=None, port=None, pool_size=None,
           - `port` (optional): port number on which to connect
           - `pool_size` (optional): DEPRECATED
           - `auto_start_request` (optional): DEPRECATED
-          - `slave_okay` (optional): is it okay to connect directly to and
-            perform queries on a slave instance
+          - `slave_okay` (optional): is it okay to connect directly to
+            and perform queries on a slave instance
           - `timeout` (optional): DEPRECATED
-          - `network_timeout` (optional): timeout (in seconds) to use for socket
-            operations - default is no timeout
+          - `network_timeout` (optional): timeout (in seconds) to use
+            for socket operations - default is no timeout
+          - `document_class` (optional): default class to use for
+            documents returned from queries on this connection
 
         .. seealso:: :meth:`end_request`
+        .. versionadded:: 1.6+
+           The `document_class` parameter.
         .. versionchanged:: 1.4
            DEPRECATED The `pool_size`, `auto_start_request`, and `timeout`
            parameters.
@@ -177,6 +181,7 @@ def __init__(self, host=None, port=None, pool_size=None,
         self.__pool = Pool(self.__connect)
 
         self.__network_timeout = network_timeout
+        self.__document_class = document_class
 
         # cache of existing indexes used by ensure_index ops
         self.__index_cache = {}
@@ -404,6 +409,19 @@ def slave_okay(self):
         """
         return self.__slave_okay
 
+    def get_document_class(self):
+        return self.__document_class
+
+    def set_document_class(self, klass):
+        self.__document_class = klass
+
+    document_class = property(get_document_class, set_document_class,
+                              doc="""Default class to use for documents
+                              returned from queries on this connection.
+
+                              .. versionadded:: 1.6+
+                              """)
+
     def __find_master(self):
         """Create a new socket and use it to figure out who the master is.
 
diff --git a/pymongo/cursor.py b/pymongo/cursor.py
@@ -75,8 +75,7 @@ def __init__(self, collection, spec=None, fields=None, skip=0, limit=0,
                 fields = helpers._fields_list_to_dict(fields)
 
         if as_class is None:
-            # TODO add Connection level default for this case
-            as_class = dict
+            as_class = collection.database.connection.document_class
 
         self.__collection = collection
         self.__spec = spec
diff --git a/test/test_connection.py b/test/test_connection.py
@@ -30,6 +30,7 @@
                             InvalidName,
                             InvalidURI,
                             OperationFailure)
+from pymongo.son import SON
 from test import version
 
 
@@ -347,6 +348,34 @@ def loop(pipe):
         except EOFError:
             pass
 
+    def test_document_class(self):
+        c = Connection(self.host, self.port)
+        db = c.pymongo_test
+        db.test.insert({"x": 1})
+
+        self.assertEqual(dict, c.document_class)
+        self.assert_(isinstance(db.test.find_one(), dict))
+        self.failIf(isinstance(db.test.find_one(), SON))
+
+        c.document_class = SON
+
+        self.assertEqual(SON, c.document_class)
+        self.assert_(isinstance(db.test.find_one(), SON))
+        self.failIf(isinstance(db.test.find_one(as_class=dict), SON))
+
+        c = Connection(self.host, self.port, document_class=SON)
+        db = c.pymongo_test
+
+        self.assertEqual(SON, c.document_class)
+        self.assert_(isinstance(db.test.find_one(), SON))
+        self.failIf(isinstance(db.test.find_one(as_class=dict), SON))
+
+        c.document_class = dict
+
+        self.assertEqual(dict, c.document_class)
+        self.assert_(isinstance(db.test.find_one(), dict))
+        self.failIf(isinstance(db.test.find_one(), SON))
+
 # TODO come up with a different way to test `network_timeout`. This is just
 # too sketchy.
 #
