PYTHON-5395 - Add convert_decimal to CodecOptions

NoahStapp · NoahStapp · commit b24288e9c824 · 2025-08-19T07:49:44.000-07:00
diff --git a/bson/__init__.py b/bson/__init__.py
@@ -72,6 +72,7 @@
 from __future__ import annotations
 
 import datetime
+import decimal
 import itertools
 import os
 import re
@@ -858,6 +859,16 @@ def _encode_decimal128(name: bytes, value: Decimal128, dummy0: Any, dummy1: Any)
     return b"\x13" + name + value.bid
 
 
+def _encode_python_decimal(
+    name: bytes, value: decimal.Decimal, dummy0: Any, opts: CodecOptions[Any]
+) -> bytes:
+    if opts.convert_decimal:
+        converted = Decimal128(value)
+        return b"\x13" + name + converted.bid
+    else:
+        raise InvalidDocument("decimal.Decimal must be converted to bson.decimal128.Decimal128.")
+
+
 def _encode_minkey(name: bytes, dummy0: Any, dummy1: Any, dummy2: Any) -> bytes:
     """Encode bson.min_key.MinKey."""
     return b"\xFF" + name
@@ -885,6 +896,7 @@ def _encode_maxkey(name: bytes, dummy0: Any, dummy1: Any, dummy2: Any) -> bytes:
     str: _encode_text,
     tuple: _encode_list,
     type(None): _encode_none,
+    decimal.Decimal: _encode_python_decimal,
     uuid.UUID: _encode_uuid,
     Binary: _encode_binary,
     Int64: _encode_long,
diff --git a/bson/_cbsonmodule.c b/bson/_cbsonmodule.c
@@ -159,6 +159,27 @@ extern int cbson_long_long_to_str(long long num, char* str, size_t size) {
     return 0;
 }
 
+int _check_decimal(PyObject *value) {
+    static PyObject *decimal_module = NULL;
+    static PyObject *decimal_class = NULL;
+
+    if (decimal_module == NULL) {
+        decimal_module = PyImport_ImportModule("decimal");
+        if (decimal_module == NULL) {
+            PyErr_SetString(PyExc_ImportError, "Failed to import decimal module");
+            return -1;
+        }
+        decimal_class = PyObject_GetAttrString(decimal_module, "Decimal");
+        if (decimal_class == NULL) {
+            Py_DECREF(decimal_module);
+            decimal_module = NULL;
+            PyErr_SetString(PyExc_AttributeError, "Failed to get Decimal class");
+            return -1;
+        }
+    }
+    return PyObject_IsInstance(value, decimal_class);
+}
+
 static PyObject* _test_long_long_to_str(PyObject* self, PyObject* args) {
     // Test extreme values
     Py_ssize_t maxNum = PY_SSIZE_T_MAX;
@@ -791,14 +812,15 @@ int convert_codec_options(PyObject* self, PyObject* options_obj, codec_options_t
 
     options->unicode_decode_error_handler = NULL;
 
-    if (!PyArg_ParseTuple(options_obj, "ObbzOOb",
+    if (!PyArg_ParseTuple(options_obj, "ObbzOObb",
                           &options->document_class,
                           &options->tz_aware,
                           &options->uuid_rep,
                           &options->unicode_decode_error_handler,
                           &options->tzinfo,
                           &type_registry_obj,
-                          &options->datetime_conversion)) {
+                          &options->datetime_conversion,
+                          &options->convert_decimal)) {
         return 0;
     }
 
@@ -993,6 +1015,26 @@ static int _write_regex_to_buffer(
     return 1;
 }
 
+static int _write_decimal_128_to_buffer(struct module_state *state, PyObject* value, buffer_t buffer, int type_byte) {
+    const char* data;
+    PyObject* pystring = PyObject_GetAttr(value, state->_bid_str);
+    if (!pystring) {
+        return 0;
+    }
+    data = PyBytes_AsString(pystring);
+    if (!data) {
+        Py_DECREF(pystring);
+        return 0;
+    }
+    if (!buffer_write_bytes(buffer, data, 16)) {
+        Py_DECREF(pystring);
+        return 0;
+    }
+    Py_DECREF(pystring);
+    *(pymongo_buffer_get_buffer(buffer) + type_byte) = 0x13;
+    return 1;
+}
+
 /* Write a single value to the buffer (also write its type_byte, for which
  * space has already been reserved.
  *
@@ -1206,23 +1248,7 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
     case 19:
         {
             /* Decimal128 */
-            const char* data;
-            PyObject* pystring = PyObject_GetAttr(value, state->_bid_str);
-            if (!pystring) {
-                return 0;
-            }
-            data = PyBytes_AsString(pystring);
-            if (!data) {
-                Py_DECREF(pystring);
-                return 0;
-            }
-            if (!buffer_write_bytes(buffer, data, 16)) {
-                Py_DECREF(pystring);
-                return 0;
-            }
-            Py_DECREF(pystring);
-            *(pymongo_buffer_get_buffer(buffer) + type_byte) = 0x13;
-            return 1;
+            return _write_decimal_128_to_buffer(state, value, buffer, type_byte);
         }
     case 100:
         {
@@ -1436,6 +1462,16 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
                                           in_fallback_call);
         Py_DECREF(binary_value);
         return result;
+    } else if (options->convert_decimal && _check_decimal(value)) {
+            /* Convert decimal.Decimal to Decimal128 */
+            PyObject* args = PyTuple_New(1);
+
+            Py_INCREF(value);
+            PyTuple_SetItem(args, 0, value);
+            PyObject* converted = PyObject_CallObject(state->Decimal128, args);
+            Py_DECREF(args);
+
+            return _write_decimal_128_to_buffer(state, converted, buffer, type_byte);
     }
 
     /* Try a custom encoder if one is provided and we have not already
diff --git a/bson/_cbsonmodule.h b/bson/_cbsonmodule.h
@@ -70,6 +70,7 @@ typedef struct codec_options_t {
     PyObject* tzinfo;
     type_registry_t type_registry;
     unsigned char datetime_conversion;
+    unsigned char convert_decimal;
     PyObject* options_obj;
     unsigned char is_raw_bson;
 } codec_options_t;
diff --git a/bson/codec_options.py b/bson/codec_options.py
@@ -241,6 +241,7 @@ class _BaseCodecOptions(NamedTuple):
     tzinfo: Optional[datetime.tzinfo]
     type_registry: TypeRegistry
     datetime_conversion: Optional[DatetimeConversion]
+    convert_decimal: Optional[bool]
 
 
 if TYPE_CHECKING:
@@ -253,6 +254,7 @@ class CodecOptions(Tuple[_DocumentType], Generic[_DocumentType]):
         tzinfo: Optional[datetime.tzinfo]
         type_registry: TypeRegistry
         datetime_conversion: Optional[int]
+        convert_decimal: Optional[bool]
 
         def __new__(
             cls: Type[CodecOptions[_DocumentType]],
@@ -263,6 +265,7 @@ def __new__(
             tzinfo: Optional[datetime.tzinfo] = ...,
             type_registry: Optional[TypeRegistry] = ...,
             datetime_conversion: Optional[int] = ...,
+            convert_decimal: Optional[bool] = ...,
         ) -> CodecOptions[_DocumentType]:
             ...
 
@@ -362,6 +365,9 @@ def __init__(self, *args, **kwargs):
                 return DatetimeMS objects when the underlying datetime is
                 out-of-range and 'datetime_clamp' to clamp to the minimum and
                 maximum possible datetimes. Defaults to 'datetime'.
+            :param convert_decimal: If ``True``, instances of :class:`~decimal.Decimal` will
+            be automatically converted to :class:`~bson.decimal128.Decimal128` when encoding to BSON.
+            Defaults to ``False``.
 
             .. versionchanged:: 4.0
                The default for `uuid_representation` was changed from
@@ -388,6 +394,7 @@ def __new__(
             tzinfo: Optional[datetime.tzinfo] = None,
             type_registry: Optional[TypeRegistry] = None,
             datetime_conversion: Optional[DatetimeConversion] = DatetimeConversion.DATETIME,
+            convert_decimal: Optional[bool] = False,
         ) -> CodecOptions:
             doc_class = document_class or dict
             # issubclass can raise TypeError for generic aliases like SON[str, Any].
@@ -439,6 +446,7 @@ def __new__(
                     tzinfo,
                     type_registry,
                     datetime_conversion,
+                    convert_decimal,
                 ),
             )
 
@@ -455,14 +463,15 @@ def _arguments_repr(self) -> str:
             return (
                 "document_class={}, tz_aware={!r}, uuid_representation={}, "
                 "unicode_decode_error_handler={!r}, tzinfo={!r}, "
-                "type_registry={!r}, datetime_conversion={!s}".format(
+                "type_registry={!r}, datetime_conversion={!s}, convert_decimal={!s}".format(
                     document_class_repr,
                     self.tz_aware,
                     uuid_rep_repr,
                     self.unicode_decode_error_handler,
                     self.tzinfo,
                     self.type_registry,
                     self.datetime_conversion,
+                    self.convert_decimal,
                 )
             )
 
@@ -477,6 +486,7 @@ def _options_dict(self) -> dict[str, Any]:
                 "tzinfo": self.tzinfo,
                 "type_registry": self.type_registry,
                 "datetime_conversion": self.datetime_conversion,
+                "convert_decimal": self.convert_decimal,
             }
 
         def __repr__(self) -> str:
@@ -513,6 +523,7 @@ def _parse_codec_options(options: Any) -> CodecOptions[Any]:
         "tzinfo",
         "type_registry",
         "datetime_conversion",
+        "convert_decimal",
     }:
         if k == "uuidrepresentation":
             kwargs["uuid_representation"] = options[k]
diff --git a/test/test_bson.py b/test/test_bson.py
@@ -19,6 +19,7 @@
 import array
 import collections
 import datetime
+import decimal
 import mmap
 import os
 import pickle
@@ -1406,6 +1407,15 @@ def test_bson_encode_decode(self) -> None:
         decoded["new_field"] = 1
         self.assertTrue(decoded["_id"].generation_time)
 
+    def test_convert_decimal(self):
+        opts = CodecOptions(convert_decimal=True)
+        decimal128_doc = {"d": bson.Decimal128("1.0")}
+        decimal_doc = {"d": decimal.Decimal("1.0")}
+        decimal_128_encoded = bson.encode(decimal128_doc, codec_options=opts)
+        decimal_encoded = bson.encode(decimal_doc, codec_options=opts)
+
+        self.assertEqual(decimal_128_encoded, decimal_encoded)
+
 
 class TestDatetimeConversion(unittest.TestCase):
     def test_comps(self):
