MasterSlaveConnection.document_class PYTHON-136

behackett · behackett · commit 020318251939 · 2011-11-02T13:45:29.000-07:00
Document class is now configurable as in Connection
and ReplicaSetConnection.
diff --git a/pymongo/master_slave_connection.py b/pymongo/master_slave_connection.py
@@ -32,7 +32,7 @@ class MasterSlaveConnection(BaseObject):
     """A master-slave connection to Mongo.
     """
 
-    def __init__(self, master, slaves=[], tz_aware=False):
+    def __init__(self, master, slaves=[], document_class=dict, tz_aware=False):
         """Create a new Master-Slave connection.
 
         The resultant connection should be interacted with using the same
@@ -50,6 +50,8 @@ def __init__(self, master, slaves=[], tz_aware=False):
           - `master`: `Connection` instance for the writable Master
           - `slaves` (optional): list of `Connection` instances for the
             read-only slaves
+          - `document_class` (optional): default class to use for
+            documents returned from queries on this connection
           - `tz_aware` (optional): if ``True``,
             :class:`~datetime.datetime` instances returned as values
             in a document by this :class:`MasterSlaveConnection` will be timezone
@@ -73,6 +75,7 @@ def __init__(self, master, slaves=[], tz_aware=False):
         self.__in_request = False
         self.__master = master
         self.__slaves = slaves
+        self.__document_class = document_class
         self.__tz_aware = tz_aware
 
     @property
@@ -83,10 +86,15 @@ def master(self):
     def slaves(self):
         return self.__slaves
 
-    # TODO this is a temporary hack PYTHON-136 is the right solution for this
-    @property
-    def document_class(self):
-        return dict
+    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 on this connection.""")
 
     @property
     def tz_aware(self):
diff --git a/test/test_master_slave_connection.py b/test/test_master_slave_connection.py
@@ -23,6 +23,7 @@
 
 from nose.plugins.skip import SkipTest
 
+from bson.son import SON
 from bson.tz_util import utc
 from pymongo import ReadPreference
 from pymongo.errors import ConnectionFailure, InvalidName
@@ -371,6 +372,35 @@ def test_base_object(self):
         self.assertFalse(self.connection.safe)
         self.assertEqual({}, self.connection.get_lasterror_options())
 
+    def test_document_class(self):
+        c = MasterSlaveConnection(self.master, self.slaves)
+        db = c.pymongo_test
+        db.test.insert({"x": 1})
+        time.sleep(1)
+
+        self.assertEqual(dict, c.document_class)
+        self.assert_(isinstance(db.test.find_one(), dict))
+        self.assertFalse(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.assertFalse(isinstance(db.test.find_one(as_class=dict), SON))
+
+        c = MasterSlaveConnection(self.master, self.slaves, document_class=SON)
+        db = c.pymongo_test
+
+        self.assertEqual(SON, c.document_class)
+        self.assert_(isinstance(db.test.find_one(), SON))
+        self.assertFalse(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.assertFalse(isinstance(db.test.find_one(), SON))
+
     def test_tz_aware(self):
         dt = datetime.datetime.utcnow()
         conn = MasterSlaveConnection(self.master, self.slaves)
