Merge pull request twisted#177 from IlyaSkriblovsky/codec_options

support for codec_options

docs/source/NEWS.rst

@@ -10,6 +10,9 @@ Features
 - ``Collection.bulk_write()`` that maches behavior of corresponding PyMongo's method. It accepts
   an iterable of ``InsertOne``, ``UpdateOne``, ... from ``pymongo.operations``, packs them into
   batches and returns aggregated response from MongoDB.
+- ``codec_options`` properties for ``ConnectionPool``, ``Database`` and ``Collection``.
+  ``Collection.with_options(codec_options=CodecOptions(document_class=...))`` is now preferred
+  over ``Collection.find(..., as_class=...)``.
 
 Release 16.1.0 (2016-06-15)
 ---------------------------

tests/test_codec_options.py

@@ -0,0 +1,71 @@
+# coding: utf-8
+# Copyright 2015 Ilya Skriblovsky <ilyaskriblovsky@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from __future__ import absolute_import, division
+
+import datetime
+from bson import SON
+from bson.codec_options import CodecOptions
+from twisted.trial import unittest
+from twisted.internet import defer
+from txmongo.connection import MongoConnection
+from txmongo.database import Database
+
+mongo_host = "localhost"
+mongo_port = 27017
+
+
+class TestCodecOptions(unittest.TestCase):
+
+    @defer.inlineCallbacks
+    def setUp(self):
+        self.conn = MongoConnection(mongo_host, mongo_port)
+        self.db = self.conn.db
+        self.coll = self.db.coll
+        yield self.coll.insert_one({'x': 42, 'y': datetime.datetime.now()})
+
+    @defer.inlineCallbacks
+    def tearDown(self):
+        yield self.coll.drop()
+        yield self.conn.disconnect()
+
+    @defer.inlineCallbacks
+    def test_Levels(self):
+        as_son = CodecOptions(document_class=SON)
+
+        doc = yield self.coll.find_one()
+        self.assertIsInstance(doc, dict)
+
+        try:
+            conn = MongoConnection(mongo_host, mongo_port, codec_options=as_son)
+            doc = yield conn.db.coll.find_one()
+            self.assertIsInstance(doc, SON)
+        finally:
+            yield conn.disconnect()
+
+        doc = yield Database(self.conn, "db", codec_options=as_son).coll.find_one()
+        self.assertIsInstance(doc, SON)
+
+        doc = yield self.coll.with_options(codec_options=as_son).find_one()
+        self.assertIsInstance(doc, SON)
+
+    @defer.inlineCallbacks
+    def test_TzAware(self):
+        doc = yield self.coll.find_one()
+        self.assertIsNone(doc['y'].tzinfo, None)
+
+        tz_aware = CodecOptions(tz_aware=True)
+        doc = yield self.coll.with_options(codec_options=tz_aware).find_one()
+        self.assertIsNotNone(doc['y'].tzinfo, None)

txmongo/collection.py

@@ -40,11 +40,15 @@ class Collection(object):
         the name of the collection to get
 
     :param write_concern:
-        An instance of :class:`pymongo.write_concern.WriteConcern`.
+        An instance of :class:`~pymongo.write_concern.WriteConcern`.
         If ``None``, ``database.write_concern`` is used.
+
+    :param codec_options:
+        An instance of :class:`~bson.codec_options.CodecOptions`.
+        If ``None``, ``database.codec_options`` is used.
     """
 
-    def __init__(self, database, name, write_concern=None):
+    def __init__(self, database, name, write_concern=None, codec_options=None):
         if not isinstance(name, (bytes, unicode)):
             raise TypeError("TxMongo: name must be an instance of (bytes, unicode).")
 
@@ -63,6 +67,7 @@ def __init__(self, database, name, write_concern=None):
         self._database = database
         self._collection_name = unicode(name)
         self.__write_concern = write_concern
+        self.__codec_options = codec_options
 
     def __str__(self):
         return "%s.%s" % (str(self._database), self._collection_name)
@@ -110,6 +115,16 @@ def __call__(self, collection_name):
         """Get a sub-collection of this collection by name."""
         return self[collection_name]
 
+    @property
+    def codec_options(self):
+        """Read only access to the :class:`~bson.codec_options.CodecOptions`
+        of this instance.
+
+        Use ``coll.with_options(codec_options=CodecOptions(...))`` to change
+        codec options.
+        """
+        return self.__codec_options or self._database.codec_options
+
     @property
     def write_concern(self):
         """Read only access to the :class:`~pymongo.write_concern.WriteConcern`
@@ -125,16 +140,21 @@ def with_options(self, **kwargs):
 
         :param write_concern: *(keyword only)*
             new :class:`~pymongo.write_concern.WriteConcern` to use.
+
+        :param codec_options: *(keyword only)*
+            new :class:`~bson.codec_options.CodecOptions` to use.
         """
         # PyMongo's method gets several positional arguments. We support
         # only write_concern for now which is the 3rd positional argument.
         # So we are using **kwargs here to force user's code to specify
         # write_concern as named argument, so adding other args in future
         # won't break compatibility
         write_concern = kwargs.get("write_concern") or self.__write_concern
+        codec_options = kwargs.get("codec_options") or self.codec_options
 
         return Collection(self._database, self._collection_name,
-                          write_concern=write_concern)
+                          write_concern=write_concern,
+                          codec_options=codec_options)
 
     @staticmethod
     def _normalize_fields_projection(fields):
@@ -281,7 +301,7 @@ def find_with_cursor(self, spec=None, skip=0, limit=0, fields=None, filter=None,
 
         spec = self.__apply_find_filter(spec, filter)
 
-        as_class = kwargs.get("as_class", dict)
+        as_class = kwargs.get("as_class")
         proto = self._database.connection.getprotocol()
 
         def after_connection(protocol):
@@ -308,7 +328,9 @@ def after_reply(reply, protocol, this_func, fetched=0):
                 docs_count = min(docs_count, limit - fetched)
             fetched += docs_count
 
-            options = bson.codec_options.CodecOptions(document_class=as_class)
+            options = self.codec_options
+            if as_class is not None:
+                options = options._replace(document_class=as_class)
             out = [document.decode(codec_options=options) for document in documents[:docs_count]]
 
             if reply.cursor_id:

txmongo/connection.py

@@ -4,6 +4,7 @@
 
 from __future__ import absolute_import, division
 
+from bson.codec_options import DEFAULT_CODEC_OPTIONS
 from pymongo.errors import AutoReconnect, ConfigurationError, OperationFailure
 from pymongo.uri_parser import parse_uri
 from pymongo.read_preferences import ReadPreference
@@ -271,6 +272,8 @@ def __init__(self, uri="mongodb://127.0.0.1:27017", pool_size=1, ssl_context_fac
         wc_options = dict((k, v) for k, v in wc_options.items() if k in self.__wc_possible_options)
         self.__write_concern = WriteConcern(**wc_options)
 
+        self.__codec_options = kwargs.get('codec_options', DEFAULT_CODEC_OPTIONS)
+
         retry_delay = kwargs.get('retry_delay', 1.0)
         max_delay = kwargs.get('max_delay', 60.0)
         self.__pool_size = pool_size
@@ -298,6 +301,10 @@ def __init__(self, uri="mongodb://127.0.0.1:27017", pool_size=1, ssl_context_fac
     def write_concern(self):
         return self.__write_concern
 
+    @property
+    def codec_options(self):
+        return self.__codec_options
+
     def getprotocols(self):
         return self.__pool
 

txmongo/database.py

@@ -13,9 +13,11 @@
 class Database(object):
     __factory = None
 
-    def __init__(self, factory, database_name, write_concern=None):
+    def __init__(self, factory, database_name, write_concern=None,
+                 codec_options=None):
         self.__factory = factory
         self.__write_concern = write_concern
+        self.__codec_options = codec_options
         self._database_name = unicode(database_name)
 
     def __str__(self):
@@ -25,7 +27,8 @@ def __repr__(self):
         return "Database(%r, %r)" % (self.__factory, self._database_name,)
 
     def __call__(self, database_name):
-        return Database(self.__factory, database_name, self.__write_concern)
+        return Database(self.__factory, database_name, self.__write_concern,
+                        self.__codec_options)
 
     def __getitem__(self, collection_name):
         return Collection(self, collection_name)
@@ -45,6 +48,10 @@ def connection(self):
     def write_concern(self):
         return self.__write_concern or self.__factory.write_concern
 
+    @property
+    def codec_options(self):
+        return self.__codec_options or self.__factory.codec_options
+
     @timeout
     @defer.inlineCallbacks
     def command(self, command, value=1, check=True, allowable_errors=None, **kwargs):