1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
|
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
import datetime
import enum
import pytest
import orjson
class StrEnum(str, enum.Enum):
AAA = "aaa"
class IntEnum(int, enum.Enum):
ONE = 1
class IntEnumEnum(enum.IntEnum):
ONE = 1
class IntFlagEnum(enum.IntFlag):
ONE = 1
class FlagEnum(enum.Flag):
ONE = 1
class AutoEnum(enum.auto):
A = "a"
class FloatEnum(float, enum.Enum):
ONE = 1.1
class Custom:
def __init__(self, val):
self.val = val
def default(obj):
if isinstance(obj, Custom):
return obj.val
raise TypeError
class UnspecifiedEnum(enum.Enum):
A = "a"
B = 1
C = FloatEnum.ONE
D = {"d": IntEnum.ONE}
E = Custom("c")
F = datetime.datetime(1970, 1, 1)
class TestEnum:
def test_cannot_subclass(self):
"""
enum.Enum cannot be subclassed
obj->ob_type->ob_base will always be enum.EnumMeta
"""
with pytest.raises(TypeError):
class Subclass(StrEnum): # type: ignore
B = "b"
def test_arbitrary_enum(self):
assert orjson.dumps(UnspecifiedEnum.A) == b'"a"'
assert orjson.dumps(UnspecifiedEnum.B) == b"1"
assert orjson.dumps(UnspecifiedEnum.C) == b"1.1"
assert orjson.dumps(UnspecifiedEnum.D) == b'{"d":1}'
def test_custom_enum(self):
assert orjson.dumps(UnspecifiedEnum.E, default=default) == b'"c"'
def test_enum_options(self):
assert (
orjson.dumps(UnspecifiedEnum.F, option=orjson.OPT_NAIVE_UTC)
== b'"1970-01-01T00:00:00+00:00"'
)
def test_int_enum(self):
assert orjson.dumps(IntEnum.ONE) == b"1"
def test_intenum_enum(self):
assert orjson.dumps(IntEnumEnum.ONE) == b"1"
def test_intflag_enum(self):
assert orjson.dumps(IntFlagEnum.ONE) == b"1"
def test_flag_enum(self):
assert orjson.dumps(FlagEnum.ONE) == b"1"
def test_auto_enum(self):
assert orjson.dumps(AutoEnum.A) == b'"a"'
def test_float_enum(self):
assert orjson.dumps(FloatEnum.ONE) == b"1.1"
def test_str_enum(self):
assert orjson.dumps(StrEnum.AAA) == b'"aaa"'
def test_bool_enum(self):
with pytest.raises(TypeError):
class BoolEnum(bool, enum.Enum): # type: ignore
TRUE = True
def test_non_str_keys_enum(self):
assert (
orjson.dumps({StrEnum.AAA: 1}, option=orjson.OPT_NON_STR_KEYS)
== b'{"aaa":1}'
)
assert (
orjson.dumps({IntEnum.ONE: 1}, option=orjson.OPT_NON_STR_KEYS) == b'{"1":1}'
)
|
