PYTHON-1782 Allow MongoClient to be initialized with type_registry

(cherry picked from commit 599e2d7)

Author: prashantmital

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
@@ -422,6 +422,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):
@@ -576,6 +584,7 @@ def validate_tzinfo(dummy, value):
 
 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,

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
@@ -95,6 +95,7 @@ def __init__(
  host=None,
  port=None,
  document_class=dict,
+ type_registry=None,
  tz_aware=None,
  connect=None,
  **kwargs):
@@ -187,6 +188,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
@@ -421,6 +425,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.
@@ -530,6 +535,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:

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,