PYTHON-5061 - Add an API to extend the bson TypeRegistry

NoahStapp · NoahStapp · commit fee5fbf33eca · 2025-05-20T10:41:07.000-04:00
diff --git a/bson/codec_options.py b/bson/codec_options.py
@@ -108,8 +108,7 @@ class TypeCodec(TypeEncoder, TypeDecoder):
 
 class TypeRegistry:
     """Encapsulates type codecs used in encoding and / or decoding BSON, as
-    well as the fallback encoder. Type registries cannot be modified after
-    instantiation.
+    well as the fallback encoder.
 
     ``TypeRegistry`` can be initialized with an iterable of type codecs, and
     a callable for the fallback encoder::
@@ -147,18 +146,33 @@ def __init__(
                 raise TypeError("fallback_encoder %r is not a callable" % (fallback_encoder))
 
         for codec in self.__type_codecs:
-            is_valid_codec = False
-            if isinstance(codec, TypeEncoder):
-                self._validate_type_encoder(codec)
-                is_valid_codec = True
-                self._encoder_map[codec.python_type] = codec.transform_python
-            if isinstance(codec, TypeDecoder):
-                is_valid_codec = True
-                self._decoder_map[codec.bson_type] = codec.transform_bson
-            if not is_valid_codec:
-                raise TypeError(
-                    f"Expected an instance of {TypeEncoder.__name__}, {TypeDecoder.__name__}, or {TypeCodec.__name__}, got {codec!r} instead"
-                )
+            self.add_codec(codec)
+
+    def add_codec(self, codec: _Codec) -> None:
+        """Adds a type codec to the registry.
+
+        :param codec: A type codec instance.
+          If a codec already exists that transforms the same python or BSON type,
+          it will be overwritten by ``codec``.
+          A TypeError will be raised if ``codec`` modifies the encoding behavior
+          of a built-in :mod:`bson` type.
+
+        .. versionadded:: 4.14
+        """
+        is_valid_codec = False
+        if isinstance(codec, TypeEncoder):
+            self._validate_type_encoder(codec)
+            is_valid_codec = True
+            self._encoder_map[codec.python_type] = codec.transform_python
+        if isinstance(codec, TypeDecoder):
+            is_valid_codec = True
+            self._decoder_map[codec.bson_type] = codec.transform_bson
+        if is_valid_codec and codec not in self.__type_codecs:
+            self.__type_codecs.append(codec)
+        elif not is_valid_codec:
+            raise TypeError(
+                f"Expected an instance of {TypeEncoder.__name__}, {TypeDecoder.__name__}, or {TypeCodec.__name__}, got {codec!r} instead"
+            )
 
     def _validate_type_encoder(self, codec: _Codec) -> None:
         from bson import _BUILT_IN_TYPES
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -1,6 +1,14 @@
 Changelog
 =========
 
+Changes in Version 4.14.0 (XXXX/XX/XX)
+--------------------------------------
+PyMongo 4.14 brings a number of changes including:
+
+- Added :meth:`bson.codec_options.TypeRegistry.add_codec` to allow users to extend :class:`bson.codec_options.TypeRegistry`
+  instances with additional type codecs after instantiation.
+
+
 Changes in Version 4.13.0 (2025/05/14)
 --------------------------------------
 
diff --git a/test/asynchronous/test_custom_types.py b/test/asynchronous/test_custom_types.py
@@ -592,46 +592,99 @@ def test_type_registry_eq(self):
         codec_instances_2 = [codec() for codec in self.codecs]
         self.assertNotEqual(TypeRegistry(codec_instances), TypeRegistry(codec_instances_2))
 
-    def test_builtin_types_override_fails(self):
-        def run_test(base, attrs):
-            msg = (
-                r"TypeEncoders cannot change how built-in types "
-                r"are encoded \(encoder .* transforms type .*\)"
-            )
-            for pytype in _BUILT_IN_TYPES:
-                attrs.update({"python_type": pytype, "transform_python": lambda x: x})
-                codec = type("testcodec", (base,), attrs)
-                codec_instance = codec()
+    def _run_builtin_types_override_fails_test(self, base, attrs, add_codec):
+        msg = (
+            r"TypeEncoders cannot change how built-in types "
+            r"are encoded \(encoder .* transforms type .*\)"
+        )
+        for pytype in _BUILT_IN_TYPES:
+            attrs.update({"python_type": pytype, "transform_python": lambda x: x})
+            codec = type("testcodec", (base,), attrs)
+            codec_instance = codec()
+            with self.assertRaisesRegex(TypeError, msg):
+                TypeRegistry(
+                    [
+                        codec_instance,
+                    ]
+                )
+
+            # Test only some subtypes as not all can be subclassed.
+            if pytype in [
+                bool,
+                type(None),
+                RE_TYPE,
+            ]:
+                continue
+
+            class MyType(pytype):  # type: ignore
+                pass
+
+            attrs.update({"python_type": MyType, "transform_python": lambda x: x})
+            codec = type("testcodec", (base,), attrs)
+            codec_instance = codec()
+            if add_codec:
+                type_registry = TypeRegistry()
+                with self.assertRaisesRegex(TypeError, msg):
+                    type_registry.add_codec(codec_instance)
+            else:
                 with self.assertRaisesRegex(TypeError, msg):
                     TypeRegistry(
                         [
                             codec_instance,
                         ]
                     )
 
-                # Test only some subtypes as not all can be subclassed.
-                if pytype in [
-                    bool,
-                    type(None),
-                    RE_TYPE,
-                ]:
-                    continue
+    def test_builtin_types_override_fails(self):
+        self._run_builtin_types_override_fails_test(TypeEncoder, {}, add_codec=False)
+        self._run_builtin_types_override_fails_test(
+            TypeCodec, {"bson_type": Decimal128, "transform_bson": lambda x: x}, add_codec=False
+        )
 
-                class MyType(pytype):  # type: ignore
-                    pass
+    def test_add_codec(self):
+        codec_instances = [codec() for codec in self.codecs]
 
-                attrs.update({"python_type": MyType, "transform_python": lambda x: x})
-                codec = type("testcodec", (base,), attrs)
-                codec_instance = codec()
-                with self.assertRaisesRegex(TypeError, msg):
-                    TypeRegistry(
-                        [
-                            codec_instance,
-                        ]
-                    )
+        def assert_proper_initialization(type_registry, codec_instances):
+            self.assertEqual(
+                type_registry._encoder_map,
+                {
+                    self.types[0]: codec_instances[0].transform_python,
+                },
+            )
+            self.assertEqual(
+                type_registry._decoder_map,
+                {int: codec_instances[0].transform_bson},
+            )
+
+        type_registry = TypeRegistry(codec_instances[:1])
+        assert_proper_initialization(type_registry, codec_instances)
 
-        run_test(TypeEncoder, {})
-        run_test(TypeCodec, {"bson_type": Decimal128, "transform_bson": lambda x: x})
+        type_registry.add_codec(codec_instances[1])
+        self.assertEqual(
+            type_registry._encoder_map,
+            {
+                self.types[0]: codec_instances[0].transform_python,
+                self.types[1]: codec_instances[1].transform_python,
+            },
+        )
+        self.assertEqual(
+            type_registry._decoder_map,
+            {int: codec_instances[0].transform_bson, str: codec_instances[1].transform_bson},
+        )
+
+    def test_add_codec_fails(self):
+        err_msg = "Expected an instance of TypeEncoder, TypeDecoder, or TypeCodec, got .* instead"
+        type_registry = TypeRegistry([codec() for codec in self.codecs])
+        with self.assertRaisesRegex(TypeError, err_msg):
+            type_registry.add_codec(self.codecs[0])  # type: ignore[arg-type]
+
+        with self.assertRaisesRegex(TypeError, err_msg):
+            type_registry.add_codec(str)  # type: ignore[arg-type]
+
+    def test_builtin_types_override_add_codec_fails(self):
+        self._run_builtin_types_override_fails_test(TypeEncoder, {}, add_codec=True)
+        self._run_builtin_types_override_fails_test(
+            TypeCodec, {"bson_type": Decimal128, "transform_bson": lambda x: x}, add_codec=True
+        )
 
 
 class TestCollectionWCustomType(AsyncIntegrationTest):
diff --git a/test/test_custom_types.py b/test/test_custom_types.py
@@ -592,46 +592,99 @@ def test_type_registry_eq(self):
         codec_instances_2 = [codec() for codec in self.codecs]
         self.assertNotEqual(TypeRegistry(codec_instances), TypeRegistry(codec_instances_2))
 
-    def test_builtin_types_override_fails(self):
-        def run_test(base, attrs):
-            msg = (
-                r"TypeEncoders cannot change how built-in types "
-                r"are encoded \(encoder .* transforms type .*\)"
-            )
-            for pytype in _BUILT_IN_TYPES:
-                attrs.update({"python_type": pytype, "transform_python": lambda x: x})
-                codec = type("testcodec", (base,), attrs)
-                codec_instance = codec()
+    def _run_builtin_types_override_fails_test(self, base, attrs, add_codec):
+        msg = (
+            r"TypeEncoders cannot change how built-in types "
+            r"are encoded \(encoder .* transforms type .*\)"
+        )
+        for pytype in _BUILT_IN_TYPES:
+            attrs.update({"python_type": pytype, "transform_python": lambda x: x})
+            codec = type("testcodec", (base,), attrs)
+            codec_instance = codec()
+            with self.assertRaisesRegex(TypeError, msg):
+                TypeRegistry(
+                    [
+                        codec_instance,
+                    ]
+                )
+
+            # Test only some subtypes as not all can be subclassed.
+            if pytype in [
+                bool,
+                type(None),
+                RE_TYPE,
+            ]:
+                continue
+
+            class MyType(pytype):  # type: ignore
+                pass
+
+            attrs.update({"python_type": MyType, "transform_python": lambda x: x})
+            codec = type("testcodec", (base,), attrs)
+            codec_instance = codec()
+            if add_codec:
+                type_registry = TypeRegistry()
+                with self.assertRaisesRegex(TypeError, msg):
+                    type_registry.add_codec(codec_instance)
+            else:
                 with self.assertRaisesRegex(TypeError, msg):
                     TypeRegistry(
                         [
                             codec_instance,
                         ]
                     )
 
-                # Test only some subtypes as not all can be subclassed.
-                if pytype in [
-                    bool,
-                    type(None),
-                    RE_TYPE,
-                ]:
-                    continue
+    def test_builtin_types_override_fails(self):
+        self._run_builtin_types_override_fails_test(TypeEncoder, {}, add_codec=False)
+        self._run_builtin_types_override_fails_test(
+            TypeCodec, {"bson_type": Decimal128, "transform_bson": lambda x: x}, add_codec=False
+        )
 
-                class MyType(pytype):  # type: ignore
-                    pass
+    def test_add_codec(self):
+        codec_instances = [codec() for codec in self.codecs]
 
-                attrs.update({"python_type": MyType, "transform_python": lambda x: x})
-                codec = type("testcodec", (base,), attrs)
-                codec_instance = codec()
-                with self.assertRaisesRegex(TypeError, msg):
-                    TypeRegistry(
-                        [
-                            codec_instance,
-                        ]
-                    )
+        def assert_proper_initialization(type_registry, codec_instances):
+            self.assertEqual(
+                type_registry._encoder_map,
+                {
+                    self.types[0]: codec_instances[0].transform_python,
+                },
+            )
+            self.assertEqual(
+                type_registry._decoder_map,
+                {int: codec_instances[0].transform_bson},
+            )
+
+        type_registry = TypeRegistry(codec_instances[:1])
+        assert_proper_initialization(type_registry, codec_instances)
 
-        run_test(TypeEncoder, {})
-        run_test(TypeCodec, {"bson_type": Decimal128, "transform_bson": lambda x: x})
+        type_registry.add_codec(codec_instances[1])
+        self.assertEqual(
+            type_registry._encoder_map,
+            {
+                self.types[0]: codec_instances[0].transform_python,
+                self.types[1]: codec_instances[1].transform_python,
+            },
+        )
+        self.assertEqual(
+            type_registry._decoder_map,
+            {int: codec_instances[0].transform_bson, str: codec_instances[1].transform_bson},
+        )
+
+    def test_add_codec_fails(self):
+        err_msg = "Expected an instance of TypeEncoder, TypeDecoder, or TypeCodec, got .* instead"
+        type_registry = TypeRegistry([codec() for codec in self.codecs])
+        with self.assertRaisesRegex(TypeError, err_msg):
+            type_registry.add_codec(self.codecs[0])  # type: ignore[arg-type]
+
+        with self.assertRaisesRegex(TypeError, err_msg):
+            type_registry.add_codec(str)  # type: ignore[arg-type]
+
+    def test_builtin_types_override_add_codec_fails(self):
+        self._run_builtin_types_override_fails_test(TypeEncoder, {}, add_codec=True)
+        self._run_builtin_types_override_fails_test(
+            TypeCodec, {"bson_type": Decimal128, "transform_bson": lambda x: x}, add_codec=True
+        )
 
 
 class TestCollectionWCustomType(IntegrationTest):
