Skip to content

Commit 4810273

Browse files
NyanFisherMaksim Zaikin
NyanFisher
authored and
Maksim Zaikin
committed
Add callable support for decimal_format option
1 parent 4024c90 commit 4810273

5 files changed

Lines changed: 184 additions & 36 deletions

File tree

‎src/msgspec/_core.c‎

Lines changed: 94 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -9460,6 +9460,7 @@ AssocList_FromObject(PyObject *obj) {
94609460
enum decimal_format {
94619461
DECIMAL_FORMAT_STRING = 0,
94629462
DECIMAL_FORMAT_NUMBER = 1,
9463+
DECIMAL_FORMAT_CALLABLE = 2,
94639464
};
94649465

94659466
enum uuid_format {
@@ -9471,6 +9472,7 @@ enum uuid_format {
94719472
typedef struct EncoderState {
94729473
MsgspecState *mod; /* module reference */
94739474
PyObject *enc_hook; /* `enc_hook` callback */
9475+
PyObject *decimal_callable; /* `decimal_format` callable */
94749476
enum decimal_format decimal_format;
94759477
enum uuid_format uuid_format;
94769478
enum order_mode order;
@@ -9485,6 +9487,7 @@ typedef struct EncoderState {
94859487
typedef struct Encoder {
94869488
PyObject_HEAD
94879489
PyObject *enc_hook;
9490+
PyObject *decimal_callable;
94889491
MsgspecState *mod;
94899492
enum decimal_format decimal_format;
94909493
enum uuid_format uuid_format;
@@ -9567,12 +9570,12 @@ Encoder_init(Encoder *self, PyObject *args, PyObject *kwds)
95679570
}
95689571

95699572
/* Process decimal format */
9573+
self->decimal_callable = NULL;
95709574
if (decimal_format == NULL) {
95719575
self->decimal_format = DECIMAL_FORMAT_STRING;
95729576
}
9573-
else {
9577+
else if (PyUnicode_CheckExact(decimal_format)) {
95749578
bool ok = false;
9575-
if (PyUnicode_CheckExact(decimal_format)) {
95769579
if (PyUnicode_CompareWithASCIIString(decimal_format, "string") == 0) {
95779580
self->decimal_format = DECIMAL_FORMAT_STRING;
95789581
ok = true;
@@ -9581,16 +9584,28 @@ Encoder_init(Encoder *self, PyObject *args, PyObject *kwds)
95819584
self->decimal_format = DECIMAL_FORMAT_NUMBER;
95829585
ok = true;
95839586
}
9584-
}
95859587
if (!ok) {
95869588
PyErr_Format(
95879589
PyExc_ValueError,
9588-
"`decimal_format` must be 'string' or 'number', got %R",
9590+
"`decimal_format` must be 'string', 'number', or a callable, got %R",
95899591
decimal_format
95909592
);
95919593
return -1;
95929594
}
95939595
}
9596+
else if (PyCallable_Check(decimal_format)) {
9597+
self->decimal_format = DECIMAL_FORMAT_CALLABLE;
9598+
Py_INCREF(decimal_format);
9599+
self->decimal_callable = decimal_format;
9600+
}
9601+
else {
9602+
PyErr_Format(
9603+
PyExc_TypeError,
9604+
"`decimal_format` must be 'string', 'number', or a callable, got %R",
9605+
decimal_format
9606+
);
9607+
return -1;
9608+
}
95949609

95959610
/* Process uuid format */
95969611
if (uuid_format == NULL) {
@@ -9637,13 +9652,15 @@ static int
96379652
Encoder_traverse(Encoder *self, visitproc visit, void *arg)
96389653
{
96399654
Py_VISIT(self->enc_hook);
9655+
Py_VISIT(self->decimal_callable);
96409656
return 0;
96419657
}
96429658

96439659
static int
96449660
Encoder_clear(Encoder *self)
96459661
{
96469662
Py_CLEAR(self->enc_hook);
9663+
Py_CLEAR(self->decimal_callable);
96479664
return 0;
96489665
}
96499666

@@ -9718,6 +9735,7 @@ encoder_encode_into_common(
97189735
.mod = self->mod,
97199736
.enc_hook = self->enc_hook,
97209737
.decimal_format = self->decimal_format,
9738+
.decimal_callable = self->decimal_callable,
97219739
.uuid_format = self->uuid_format,
97229740
.order = self->order,
97239741
.output_buffer = buf,
@@ -9765,6 +9783,7 @@ encoder_encode_common(
97659783
.mod = self->mod,
97669784
.enc_hook = self->enc_hook,
97679785
.decimal_format = self->decimal_format,
9786+
.decimal_callable = self->decimal_callable,
97689787
.uuid_format = self->uuid_format,
97699788
.order = self->order,
97709789
.output_len = 0,
@@ -9822,6 +9841,7 @@ encode_common(
98229841
.mod = mod,
98239842
.enc_hook = enc_hook,
98249843
.decimal_format = DECIMAL_FORMAT_STRING,
9844+
.decimal_callable = NULL,
98259845
.uuid_format = UUID_FORMAT_CANONICAL,
98269846
.output_len = 0,
98279847
.max_output_len = ENC_INIT_BUFSIZE,
@@ -9853,7 +9873,13 @@ Encoder_decimal_format(Encoder *self, void *closure) {
98539873
if (self->decimal_format == DECIMAL_FORMAT_STRING) {
98549874
return PyUnicode_InternFromString("string");
98559875
}
9856-
return PyUnicode_InternFromString("number");
9876+
else if (self->decimal_format == DECIMAL_FORMAT_NUMBER) {
9877+
return PyUnicode_InternFromString("number");
9878+
}
9879+
else {
9880+
Py_INCREF(self->decimal_callable);
9881+
return self->decimal_callable;
9882+
}
98579883
}
98589884

98599885
static PyObject*
@@ -12423,11 +12449,12 @@ PyDoc_STRVAR(Encoder__doc__,
1242312449
" A callable to call for objects that aren't supported msgspec types. Takes\n"
1242412450
" the unsupported object and should return a supported object, or raise a\n"
1242512451
" ``NotImplementedError`` if unsupported.\n"
12426-
"decimal_format : {'string', 'number'}, optional\n"
12452+
"decimal_format : {'string', 'number'} or callable, optional\n"
1242712453
" The format to use for encoding `decimal.Decimal` objects. If 'string'\n"
12428-
" they're encoded as strings, if 'number', they're encoded as floats.\n"
12429-
" Defaults to 'string', which is the recommended value since 'number'\n"
12430-
" may result in precision loss when decoding.\n"
12454+
" they're encoded as strings, if 'number', they're encoded as floats. If a\n"
12455+
" callable is provided, it will be called with the Decimal object and should\n"
12456+
" return a supported object. Defaults to 'string', which is the recommended\n"
12457+
" value since 'number' may result in precision loss when decoding.\n"
1243112458
"uuid_format : {'canonical', 'hex', 'bytes'}, optional\n"
1243212459
" The format to use for encoding `uuid.UUID` objects. The 'canonical'\n"
1243312460
" and 'hex' formats encode them as strings with and without hyphens\n"
@@ -13311,11 +13338,27 @@ mpack_encode_decimal(EncoderState *self, PyObject *obj)
1331113338
if (temp == NULL) return -1;
1331213339
out = mpack_encode_str(self, temp);
1331313340
}
13314-
else {
13341+
else if (self->decimal_format == DECIMAL_FORMAT_NUMBER) {
1331513342
temp = PyNumber_Float(obj);
1331613343
if (temp == NULL) return -1;
1331713344
out = mpack_encode_float(self, temp);
1331813345
}
13346+
else {
13347+
temp = PyObject_CallFunctionObjArgs(self->decimal_callable, obj, NULL);
13348+
if (temp == NULL) return -1;
13349+
13350+
PyTypeObject *type = Py_TYPE(temp);
13351+
if (type == (PyTypeObject *)(self->mod->DecimalType)) {
13352+
PyErr_SetString(
13353+
PyExc_TypeError,
13354+
"decimal_format callable must not return a Decimal"
13355+
);
13356+
Py_DECREF(temp);
13357+
return -1;
13358+
}
13359+
13360+
out = mpack_encode(self, temp);
13361+
}
1331913362
Py_DECREF(temp);
1332013363
return out;
1332113364
}
@@ -13639,12 +13682,13 @@ PyDoc_STRVAR(JSONEncoder__doc__,
1363913682
" A callable to call for objects that aren't supported msgspec types. Takes\n"
1364013683
" the unsupported object and should return a supported object, or raise a\n"
1364113684
" ``NotImplementedError`` if unsupported.\n"
13642-
"decimal_format : {'string', 'number'}, optional\n"
13685+
"decimal_format : {'string', 'number'} or callable, optional\n"
1364313686
" The format to use for encoding `decimal.Decimal` objects. If 'string'\n"
13644-
" they're encoded as strings, if 'number', they're encoded as floats.\n"
13645-
" Defaults to 'string', which is the recommended value since 'number'\n"
13646-
" may result in precision loss when decoding for some JSON library\n"
13647-
" implementations.\n"
13687+
" they're encoded as strings, if 'number', they're encoded as floats. If a\n"
13688+
" callable is provided, it will be called with the Decimal object and should\n"
13689+
" return a supported object. Defaults to 'string', which is the recommended\n"
13690+
" value since 'number' may result in precision loss when decoding for some\n"
13691+
" JSON library implementations.\n"
1364813692
"uuid_format : {'canonical', 'hex'}, optional\n"
1364913693
" The format to use for encoding `uuid.UUID` objects. The 'canonical'\n"
1365013694
" and 'hex' formats encode them as strings with and without hyphens\n"
@@ -13951,27 +13995,46 @@ json_encode_uuid(EncoderState *self, PyObject *obj)
1395113995
static int
1395213996
json_encode_decimal(EncoderState *self, PyObject *obj)
1395313997
{
13954-
PyObject *temp = PyObject_Str(obj);
13955-
if (temp == NULL) return -1;
13998+
PyObject *temp;
1395613999

13957-
Py_ssize_t size;
13958-
const char* buf = unicode_str_and_size_nocheck(temp, &size);
13959-
bool decimal_as_string = (self->decimal_format == DECIMAL_FORMAT_STRING);
14000+
if (MS_LIKELY(self->decimal_format != DECIMAL_FORMAT_CALLABLE)) {
14001+
temp = PyObject_Str(obj);
14002+
if (temp == NULL) return -1;
14003+
14004+
Py_ssize_t size;
14005+
const char* buf = unicode_str_and_size_nocheck(temp, &size);
14006+
bool decimal_as_string = (self->decimal_format == DECIMAL_FORMAT_STRING);
1396014007

13961-
Py_ssize_t required = size + (2 * decimal_as_string);
13962-
if (ms_ensure_space(self, size + 2) < 0) {
14008+
Py_ssize_t required = size + (2 * decimal_as_string);
14009+
if (ms_ensure_space(self, size + 2) < 0) {
14010+
Py_DECREF(temp);
14011+
return -1;
14012+
}
14013+
char *p = self->output_buffer_raw + self->output_len;
14014+
if (MS_LIKELY(decimal_as_string)) *p++ = '"';
14015+
memcpy(p, buf, size);
14016+
if (MS_LIKELY(decimal_as_string)) *(p + size) = '"';
14017+
self->output_len += required;
1396314018
Py_DECREF(temp);
13964-
return -1;
1396514019
}
14020+
else {
14021+
temp = PyObject_CallFunctionObjArgs(self->decimal_callable, obj, NULL);
14022+
if (temp == NULL) return -1;
1396614023

13967-
char *p = self->output_buffer_raw + self->output_len;
13968-
if (MS_LIKELY(decimal_as_string)) *p++ = '"';
13969-
memcpy(p, buf, size);
13970-
if (MS_LIKELY(decimal_as_string)) *(p + size) = '"';
13971-
13972-
self->output_len += required;
14024+
PyTypeObject *type = Py_TYPE(temp);
14025+
if (type == (PyTypeObject *)(self->mod->DecimalType)) {
14026+
PyErr_SetString(
14027+
PyExc_TypeError,
14028+
"decimal_format callable must not return a Decimal"
14029+
);
14030+
Py_DECREF(temp);
14031+
return -1;
14032+
}
1397314033

13974-
Py_DECREF(temp);
14034+
int out = json_encode(self, temp);
14035+
Py_DECREF(temp);
14036+
return out;
14037+
}
1397514038

1397614039
return 0;
1397714040
}
@@ -14669,6 +14732,7 @@ JSONEncoder_encode_lines(Encoder *self, PyObject *const *args, Py_ssize_t nargs)
1466914732
.mod = self->mod,
1467014733
.enc_hook = self->enc_hook,
1467114734
.decimal_format = self->decimal_format,
14735+
.decimal_callable = self->decimal_callable,
1467214736
.uuid_format = self->uuid_format,
1467314737
.order = self->order,
1467414738
.output_len = 0,

‎src/msgspec/json.pyi‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import decimal
12
from collections.abc import Callable, Iterable
23
from typing import (
34
Any,
@@ -23,15 +24,21 @@ schema_hook_sig = Optional[Callable[[type], dict[str, Any]]]
2324

2425
class Encoder:
2526
enc_hook: enc_hook_sig
26-
decimal_format: Literal["string", "number"]
27+
decimal_format: Union[
28+
Literal["string", "number"],
29+
Callable[[decimal.Decimal], Union[str, float]],
30+
]
2731
uuid_format: Literal["canonical", "hex"]
2832
order: Literal[None, "deterministic", "sorted"]
2933

3034
def __init__(
3135
self,
3236
*,
3337
enc_hook: enc_hook_sig = None,
34-
decimal_format: Literal["string", "number"] = "string",
38+
decimal_format: Union[
39+
Literal["string", "number"],
40+
Callable[[decimal.Decimal], Union[str, float]],
41+
] = "string",
3542
uuid_format: Literal["canonical", "hex"] = "canonical",
3643
order: Literal[None, "deterministic", "sorted"] = None,
3744
): ...

‎src/msgspec/msgpack.pyi‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import decimal
12
from typing import (
23
Any,
34
Callable,
@@ -60,14 +61,20 @@ class Decoder(Generic[T]):
6061

6162
class Encoder:
6263
enc_hook: enc_hook_sig
63-
decimal_format: Literal["string", "number"]
64+
decimal_format: Union[
65+
Literal["string", "number"],
66+
Callable[[decimal.Decimal], Union[str, float]],
67+
]
6468
uuid_format: Literal["canonical", "hex", "bytes"]
6569
order: Literal[None, "deterministic", "sorted"]
6670
def __init__(
6771
self,
6872
*,
6973
enc_hook: enc_hook_sig = None,
70-
decimal_format: Literal["string", "number"] = "string",
74+
decimal_format: Union[
75+
Literal["string", "number"],
76+
Callable[[decimal.Decimal], Union[str, float]],
77+
] = "string",
7178
uuid_format: Literal["canonical", "hex", "bytes"] = "canonical",
7279
order: Literal[None, "deterministic", "sorted"] = None,
7380
): ...

‎tests/typing/basic_typing_examples.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,6 +691,7 @@ def check_msgpack_order() -> None:
691691
def check_msgpack_Encoder_decimal_format() -> None:
692692
enc = msgspec.msgpack.Encoder(decimal_format="string")
693693
msgspec.msgpack.Encoder(decimal_format="number")
694+
msgspec.msgpack.Encoder(decimal_format=lambda x: str(x))
694695
reveal_type(enc.decimal_format) # assert "string" in typ.lower() and "number" in typ.lower()
695696

696697

@@ -871,6 +872,7 @@ def check_json_order() -> None:
871872
def check_json_Encoder_decimal_format() -> None:
872873
enc = msgspec.json.Encoder(decimal_format="string")
873874
msgspec.json.Encoder(decimal_format="number")
875+
msgspec.json.Encoder(decimal_format=lambda x: str(x))
874876
reveal_type(enc.decimal_format) # assert "string" in typ.lower() and "number" in typ.lower()
875877

876878

0 commit comments

Comments
 (0)