Skip to content

Commit 9999e90

Browse files
ijlijl
committed
datetime.datetime option to assume UTC
1 parent 5920ac0 commit 9999e90

5 files changed

Lines changed: 109 additions & 41 deletions

File tree

‎README.md‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,12 +47,22 @@ deploying this does not require Rust or non-libc type libraries.)
4747
def dumps(obj: Any, default=Optional[Callable[Any]], option=Optional[int]) -> bytes: ...
4848
```
4949

50-
`dumps()` serializes Python objects to JSON. It natively serializes
50+
`dumps()` serializes Python objects to JSON.
51+
52+
It natively serializes
5153
`str`, `dict`, `list`, `tuple`, `int`, `float`, `datetime.datetime`,
5254
`datetime.date`, `datetime.time`, and `None` instances. It supports
53-
arbitrary types through `default`. It does not serialize
55+
arbitrary types through `default`.
56+
57+
It does not serialize
5458
subclasses of supported types natively, but `default` may be used.
5559

60+
It accepts options via an `option` keyword argument. These include
61+
`orjson.OPT_STRICT_INTEGER` for enforcing a 53-bit limit on integers
62+
and `orjson.OPT_NAIVE_UTC` for assuming `datetime.datetime` objects without a
63+
`tzinfo` are UTC. Specify multiple options by masking them together, e.g.,
64+
`option=orjson.OPT_STRICT_INTEGER | orjson.OPT_NAIVE_UTC`.
65+
5666
It raises `JSONEncodeError` on an unsupported type. This exception message
5767
describes the invalid object.
5868

@@ -156,10 +166,17 @@ ISO 8601.
156166

157167
`datetime.datetime` objects must have `tzinfo` set. For UTC timezones,
158168
`datetime.timezone.utc` is sufficient. For other timezones, `tzinfo`
159-
must be a timezone object from the pendulum, pytz, or dateutil libraries.
169+
must be a timezone object from the pendulum, pytz, or dateutil libraries. For
170+
applications in which naive datetimes are known to be UTC, `tzinfo` may be
171+
omitted if `orjson.OPT_NAIVE_UTC` if specified. This does not affect
172+
datetimes with a `tzinfo` set.
160173

161174
```python
162175
>>> import orjson, datetime, pendulum
176+
>>> orjson.dumps(
177+
datetime.datetime.fromtimestamp(4123518902), option=orjson.OPT_NAIVE_UTC
178+
)
179+
b'"2100-09-01T21:55:02+00:00"'
163180
>>> orjson.dumps(
164181
datetime.datetime.fromtimestamp(4123518902).replace(tzinfo=datetime.timezone.utc)
165182
)

‎src/encode.rs‎

Lines changed: 56 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@ const STRICT_INT_MIN: i64 = -9007199254740991;
1515
const STRICT_INT_MAX: i64 = 9007199254740991;
1616

1717
pub const STRICT_INTEGER: u8 = 1 << 0;
18+
pub const NAIVE_UTC: u8 = 1 << 1;
1819

19-
pub const MAX_OPT: i8 = STRICT_INTEGER as i8;
20+
pub const MAX_OPT: i8 = STRICT_INTEGER as i8 | NAIVE_UTC as i8;
2021

2122
const HYPHEN: u8 = 45; // "-"
2223
const PLUS: u8 = 43; // "+"
@@ -177,51 +178,70 @@ impl<'p> Serialize for SerializePyObject {
177178
serializer.serialize_seq(None).unwrap().end()
178179
}
179180
} else if unsafe { obj_ptr == DATETIME_PTR } {
180-
if unsafe { (*(self.ptr as *mut pyo3::ffi::PyDateTime_DateTime)).hastzinfo == 0 } {
181+
if unsafe {
182+
(*(self.ptr as *mut pyo3::ffi::PyDateTime_DateTime)).hastzinfo == 0
183+
&& !self.opts & NAIVE_UTC == NAIVE_UTC
184+
} {
181185
return Err(ser::Error::custom(
182186
"datetime.datetime must have tzinfo set; use datetime.timezone.utc if UTC",
183187
));
184188
}
185189
let tzinfo = unsafe { pyo3::ffi::PyDateTime_DATE_GET_TZINFO(self.ptr) };
186-
let offset: *mut pyo3::ffi::PyObject;
187-
if unsafe { pyo3::ffi::PyObject_HasAttr(tzinfo, CONVERT_METHOD_STR) == 1 } {
188-
// pendulum
189-
offset = unsafe {
190-
pyo3::ffi::PyObject_CallMethodObjArgs(
191-
self.ptr,
192-
UTCOFFSET_METHOD_STR,
193-
std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
194-
)
195-
};
196-
} else if unsafe { pyo3::ffi::PyObject_HasAttr(tzinfo, NORMALIZE_METHOD_STR) == 1 }
197-
{
198-
// pytz
199-
offset = unsafe {
200-
pyo3::ffi::PyObject_CallMethodObjArgs(
190+
let offset_day: i32;
191+
let mut offset_second: i32;
192+
if unsafe { (*(self.ptr as *mut pyo3::ffi::PyDateTime_DateTime)).hastzinfo == 1 } {
193+
if unsafe { pyo3::ffi::PyObject_HasAttr(tzinfo, CONVERT_METHOD_STR) == 1 } {
194+
// pendulum
195+
let offset = unsafe {
196+
pyo3::ffi::PyObject_CallMethodObjArgs(
197+
self.ptr,
198+
UTCOFFSET_METHOD_STR,
199+
std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
200+
)
201+
};
202+
offset_second =
203+
unsafe { pyo3::ffi::PyDateTime_DELTA_GET_SECONDS(offset) as i32 };
204+
offset_day = unsafe { pyo3::ffi::PyDateTime_DELTA_GET_DAYS(offset) };
205+
} else if unsafe {
206+
pyo3::ffi::PyObject_HasAttr(tzinfo, NORMALIZE_METHOD_STR) == 1
207+
} {
208+
// pytz
209+
let offset = unsafe {
210+
pyo3::ffi::PyObject_CallMethodObjArgs(
211+
pyo3::ffi::PyObject_CallMethodObjArgs(
212+
tzinfo,
213+
NORMALIZE_METHOD_STR,
214+
self.ptr,
215+
std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
216+
),
217+
UTCOFFSET_METHOD_STR,
218+
std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
219+
)
220+
};
221+
offset_second =
222+
unsafe { pyo3::ffi::PyDateTime_DELTA_GET_SECONDS(offset) as i32 };
223+
offset_day = unsafe { pyo3::ffi::PyDateTime_DELTA_GET_DAYS(offset) };
224+
} else if unsafe { pyo3::ffi::PyObject_HasAttr(tzinfo, DST_STR) == 1 } {
225+
// dateutil/arrow, datetime.timezone.utc
226+
let offset = unsafe {
201227
pyo3::ffi::PyObject_CallMethodObjArgs(
202228
tzinfo,
203-
NORMALIZE_METHOD_STR,
229+
UTCOFFSET_METHOD_STR,
204230
self.ptr,
205231
std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
206-
),
207-
UTCOFFSET_METHOD_STR,
208-
std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
209-
)
210-
};
211-
} else if unsafe { pyo3::ffi::PyObject_HasAttr(tzinfo, DST_STR) == 1 } {
212-
// dateutil/arrow, datetime.timezone.utc
213-
offset = unsafe {
214-
pyo3::ffi::PyObject_CallMethodObjArgs(
215-
tzinfo,
216-
UTCOFFSET_METHOD_STR,
217-
self.ptr,
218-
std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
219-
)
220-
};
221-
} else {
222-
return Err(ser::Error::custom(
232+
)
233+
};
234+
offset_second =
235+
unsafe { pyo3::ffi::PyDateTime_DELTA_GET_SECONDS(offset) as i32 };
236+
offset_day = unsafe { pyo3::ffi::PyDateTime_DELTA_GET_DAYS(offset) };
237+
} else {
238+
return Err(ser::Error::custom(
223239
"datetime's timezone library is not supported: use datetime.timezone.utc, pendulum, pytz, or dateutil",
224240
));
241+
}
242+
} else {
243+
offset_second = 0;
244+
offset_day = 0;
225245
};
226246

227247
let mut dt: SmallVec<[u8; 32]> = SmallVec::with_capacity(32);
@@ -290,8 +310,6 @@ impl<'p> Serialize for SerializePyObject {
290310
dt.extend(itoa::Buffer::new().format(microsecond).bytes());
291311
}
292312
}
293-
let mut offset_second =
294-
unsafe { pyo3::ffi::PyDateTime_DELTA_GET_SECONDS(offset) as i32 };
295313
if offset_second == 0 {
296314
dt.push(PLUS);
297315
dt.push(ZERO);
@@ -300,7 +318,7 @@ impl<'p> Serialize for SerializePyObject {
300318
dt.push(ZERO);
301319
dt.push(ZERO);
302320
} else {
303-
if unsafe { pyo3::ffi::PyDateTime_DELTA_GET_DAYS(offset) == -1 } {
321+
if offset_day == -1 {
304322
// datetime.timedelta(days=-1, seconds=68400) -> -05:00
305323
dt.push(HYPHEN);
306324
offset_second = 86400 - offset_second

‎src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ fn orjson(py: Python, m: &PyModule) -> PyResult<()> {
3030
m.add("JSONDecodeError", py.get_type::<exc::JSONDecodeError>())?;
3131
m.add("JSONEncodeError", py.get_type::<exc::JSONEncodeError>())?;
3232
m.add("OPT_STRICT_INTEGER", encode::STRICT_INTEGER.into_object(py))?;
33+
m.add("OPT_NAIVE_UTC", encode::NAIVE_UTC.into_object(py))?;
3334
Ok(())
3435
}
3536

‎test/test_api.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
22

3+
import datetime
34
import unittest
45
import json
56

@@ -90,3 +91,16 @@ def test_option_range_high(self):
9091
"""
9192
with self.assertRaises(orjson.JSONEncodeError):
9293
orjson.dumps(True, option=4)
94+
95+
96+
def test_opts_multiple(self):
97+
"""
98+
dumps() multiple option
99+
"""
100+
self.assertEqual(
101+
orjson.dumps(
102+
[1, datetime.datetime.fromtimestamp(4123518902)],
103+
option=orjson.OPT_STRICT_INTEGER | orjson.OPT_NAIVE_UTC,
104+
),
105+
b'[1,"2100-09-01T21:55:02+00:00"]'
106+
)

‎test/test_datetime.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,24 @@ def test_datetime_naive(self):
1919
with self.assertRaises(TypeError):
2020
orjson.dumps([datetime.datetime(2000, 1, 1, 2, 3, 4, 123)])
2121

22+
def test_datetime_naive_utc(self):
23+
"""
24+
datetime.datetime naive with opt assumes UTC
25+
"""
26+
self.assertEqual(
27+
orjson.dumps([datetime.datetime(2000, 1, 1, 2, 3, 4, 123)], option=orjson.OPT_NAIVE_UTC),
28+
b'["2000-01-01T02:03:04.123+00:00"]'
29+
)
30+
31+
def test_datetime_tz_assume(self):
32+
"""
33+
datetime.datetime tz with assume UTC uses tz
34+
"""
35+
self.assertEqual(
36+
orjson.dumps([datetime.datetime(2018, 1, 1, 2, 3, 4, 0, tzinfo=tz.gettz('Asia/Shanghai'))], option=orjson.OPT_NAIVE_UTC),
37+
b'["2018-01-01T02:03:04+08:00"]',
38+
)
39+
2240
def test_datetime_timezone_utc(self):
2341
"""
2442
datetime.datetime UTC

0 commit comments

Comments
 (0)