20
20
# software solely pursuant to the terms of the relevant commercial agreement.
21
21
22
22
23
+ import calendar
24
+ import datetime as dt
23
25
import heapq
24
26
import io
25
27
import logging
@@ -84,19 +86,35 @@ def super_len(o):
84
86
return None
85
87
86
88
87
- def cratedb_json_encoder (obj : t .Any ) -> str :
89
+ epoch_aware = dt .datetime (1970 , 1 , 1 , tzinfo = dt .timezone .utc )
90
+ epoch_naive = dt .datetime (1970 , 1 , 1 )
91
+
92
+
93
+ def json_encoder (obj : t .Any ) -> t .Union [int , str ]:
88
94
"""
89
95
Encoder function for orjson, with additional type support.
90
96
91
- - Python's `Decimal` type.
92
- - freezegun's `FakeDatetime` type.
97
+ - Python's `Decimal` type will be serialized to `str`.
98
+ - Python's `dt.datetime` and `dt.date` types will be
99
+ serialized to `int` after converting to milliseconds
100
+ since epoch.
93
101
94
102
https://github.com/ijl/orjson#default
103
+ https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#type-timestamp
95
104
"""
96
105
if isinstance (obj , Decimal ):
97
106
return str (obj )
98
- elif hasattr (obj , "isoformat" ):
99
- return obj .isoformat ()
107
+ if isinstance (obj , dt .datetime ):
108
+ if obj .tzinfo is not None :
109
+ delta = obj - epoch_aware
110
+ else :
111
+ delta = obj - epoch_naive
112
+ return int (
113
+ delta .microseconds / 1000.0
114
+ + (delta .seconds + delta .days * 24 * 3600 ) * 1000.0
115
+ )
116
+ if isinstance (obj , dt .date ):
117
+ return calendar .timegm (obj .timetuple ()) * 1000
100
118
raise TypeError
101
119
102
120
@@ -108,8 +126,12 @@ def json_dumps(obj: t.Any) -> bytes:
108
126
"""
109
127
return orjson .dumps (
110
128
obj ,
111
- default = cratedb_json_encoder ,
112
- option = (orjson .OPT_NON_STR_KEYS | orjson .OPT_SERIALIZE_NUMPY ),
129
+ default = json_encoder ,
130
+ option = (
131
+ orjson .OPT_PASSTHROUGH_DATETIME
132
+ | orjson .OPT_NON_STR_KEYS
133
+ | orjson .OPT_SERIALIZE_NUMPY
134
+ ),
113
135
)
114
136
115
137
0 commit comments