PYTHON-1782 Allow MongoClient to be initialized with type_registry

prashantmital · prashantmital · commit 599e2d711741 · 2019-03-21T15:30:16.000-07:00
diff --git a/pymongo/common.py b/pymongo/common.py
@@ -21,7 +21,7 @@
 from bson import SON
 from bson.binary import (STANDARD, PYTHON_LEGACY,
                          JAVA_LEGACY, CSHARP_LEGACY)
-from bson.codec_options import CodecOptions
+from bson.codec_options import CodecOptions, TypeRegistry
 from bson.py3compat import abc, integer_types, iteritems, string_type
 from bson.raw_bson import RawBSONDocument
 from pymongo.auth import MECHANISMS
@@ -423,6 +423,14 @@ def validate_document_class(option, value):
     return value
 
 
+def validate_type_registry(option, value):
+    """Validate the type_registry option."""
+    if value is not None and not isinstance(value, TypeRegistry):
+        raise TypeError("%s must be an instance of %s" % (
+            option, TypeRegistry))
+    return value
+
+
 def validate_list(option, value):
     """Validates that 'value' is a list."""
     if not isinstance(value, list):
@@ -600,6 +608,7 @@ def validate_tzinfo(dummy, value):
 # values for those options.
 KW_VALIDATORS = {
     'document_class': validate_document_class,
+    'type_registry': validate_type_registry,
     'read_preference': validate_read_preference,
     'event_listeners': _validate_event_listeners,
     'tzinfo': validate_tzinfo,
diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py
@@ -39,7 +39,7 @@
 
 from collections import defaultdict
 
-from bson.codec_options import DEFAULT_CODEC_OPTIONS
+from bson.codec_options import DEFAULT_CODEC_OPTIONS, TypeRegistry
 from bson.py3compat import (integer_types,
                             string_type)
 from bson.son import SON
@@ -98,6 +98,7 @@ def __init__(
             host=None,
             port=None,
             document_class=dict,
+            type_registry=None,
             tz_aware=None,
             connect=None,
             **kwargs):
@@ -190,6 +191,9 @@ def __init__(
           - `port` (optional): port number on which to connect
           - `document_class` (optional): default class to use for
             documents returned from queries on this client
+          - `type_registry` (optional): instance of
+            :class:`~bson.codec_options.TypeRegistry` to enable encoding
+            and decoding of custom types.
           - `tz_aware` (optional): if ``True``,
             :class:`~datetime.datetime` instances returned as values
             in a document by this :class:`MongoClient` will be timezone
@@ -454,6 +458,7 @@ def __init__(
 
         .. versionchanged:: 3.8
            Added the ``server_selector`` keyword argument.
+           Added the ``type_registry`` keyword argument.
 
         .. versionchanged:: 3.7
            Added the ``driver`` keyword argument.
@@ -564,6 +569,8 @@ def __init__(
 
         keyword_opts = kwargs
         keyword_opts['document_class'] = document_class
+        if type_registry is not None:
+            keyword_opts['type_registry'] = type_registry
         if tz_aware is None:
             tz_aware = opts.get('tz_aware', False)
         if connect is None:
diff --git a/test/test_client.py b/test/test_client.py
@@ -29,7 +29,7 @@
 sys.path[0:0] = [""]
 
 from bson import BSON
-from bson.codec_options import CodecOptions
+from bson.codec_options import CodecOptions, TypeEncoder, TypeRegistry
 from bson.py3compat import thread
 from bson.son import SON
 from bson.tz_util import utc
@@ -248,14 +248,21 @@ def test_metadata(self):
         self.assertEqual(options.pool_options.metadata, metadata)
 
     def test_kwargs_codec_options(self):
+        class FloatAsIntEncoder(TypeEncoder):
+            python_type = float
+            def transform_python(self, value):
+                return int(value)
+
         # Ensure codec options are passed in correctly
         document_class = SON
+        type_registry = TypeRegistry([FloatAsIntEncoder()])
         tz_aware = True
         uuid_representation_label = 'javaLegacy'
         unicode_decode_error_handler = 'ignore'
         tzinfo = utc
         c = MongoClient(
             document_class=document_class,
+            type_registry=type_registry,
             tz_aware=tz_aware,
             uuidrepresentation=uuid_representation_label,
             unicode_decode_error_handler=unicode_decode_error_handler,
@@ -264,6 +271,7 @@ def test_kwargs_codec_options(self):
         )
 
         self.assertEqual(c.codec_options.document_class, document_class)
+        self.assertEqual(c.codec_options.type_registry, type_registry)
         self.assertEqual(c.codec_options.tz_aware, tz_aware)
         self.assertEqual(
             c.codec_options.uuid_representation,
