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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
|
import dataclasses as dc
import datetime
import decimal
import enum
import ipaddress
import operator
import pathlib
import re
import uuid
from enum import Enum
from typing import (
Any,
Callable,
Container,
Iterable,
Mapping,
MutableMapping,
Optional,
Tuple,
)
import bson
import pydantic
import beanie
from beanie.odm.fields import Link, LinkTypes
from beanie.odm.utils.pydantic import (
IS_PYDANTIC_V2,
IS_PYDANTIC_V2_10,
get_model_fields,
)
SingleArgCallable = Callable[[Any], Any]
DEFAULT_CUSTOM_ENCODERS: MutableMapping[type, SingleArgCallable] = {
ipaddress.IPv4Address: str,
ipaddress.IPv4Interface: str,
ipaddress.IPv4Network: str,
ipaddress.IPv6Address: str,
ipaddress.IPv6Interface: str,
ipaddress.IPv6Network: str,
pathlib.PurePath: str,
pydantic.SecretBytes: pydantic.SecretBytes.get_secret_value,
pydantic.SecretStr: pydantic.SecretStr.get_secret_value,
datetime.date: lambda d: datetime.datetime.combine(d, datetime.time.min),
datetime.timedelta: operator.methodcaller("total_seconds"),
enum.Enum: operator.attrgetter("value"),
Link: operator.attrgetter("ref"),
bytes: bson.Binary,
decimal.Decimal: bson.Decimal128,
uuid.UUID: bson.Binary.from_uuid,
re.Pattern: bson.Regex.from_native,
}
if IS_PYDANTIC_V2:
from pydantic_core import Url
DEFAULT_CUSTOM_ENCODERS[Url] = str
if IS_PYDANTIC_V2_10:
from pydantic import AnyUrl
DEFAULT_CUSTOM_ENCODERS[AnyUrl] = str
BSON_SCALAR_TYPES = (
type(None),
str,
int,
float,
datetime.datetime,
bson.Binary,
bson.DBRef,
bson.Decimal128,
bson.MaxKey,
bson.MinKey,
bson.ObjectId,
bson.Regex,
)
@dc.dataclass
class Encoder:
"""
BSON encoding class
"""
exclude: Container[str] = frozenset()
custom_encoders: Mapping[type, SingleArgCallable] = dc.field(
default_factory=dict
)
to_db: bool = False
keep_nulls: bool = True
def _encode_document(self, obj: "beanie.Document") -> Mapping[str, Any]:
obj.parse_store()
settings = obj.get_settings()
obj_dict = {}
if settings.union_doc is not None:
obj_dict[settings.class_id] = (
settings.union_doc_alias or obj.__class__.__name__
)
if obj._class_id:
obj_dict[settings.class_id] = obj._class_id
link_fields = obj.get_link_fields() or {}
sub_encoder = Encoder(
# don't propagate self.exclude to subdocuments
custom_encoders=settings.bson_encoders,
to_db=self.to_db,
keep_nulls=self.keep_nulls,
)
for key, value in self._iter_model_items(obj):
if key in link_fields:
link_type = link_fields[key].link_type
if link_type in (LinkTypes.DIRECT, LinkTypes.OPTIONAL_DIRECT):
if value is not None:
value = value.to_ref()
elif link_type in (LinkTypes.LIST, LinkTypes.OPTIONAL_LIST):
if value is not None:
value = [link.to_ref() for link in value]
elif self.to_db:
continue
obj_dict[key] = sub_encoder.encode(value)
return obj_dict
def encode(self, obj: Any) -> Any:
if self.custom_encoders:
encoder = _get_encoder(obj, self.custom_encoders)
if encoder is not None:
return encoder(obj)
if isinstance(obj, BSON_SCALAR_TYPES):
return obj
encoder = _get_encoder(obj, DEFAULT_CUSTOM_ENCODERS)
if encoder is not None:
return encoder(obj)
if isinstance(obj, beanie.Document):
return self._encode_document(obj)
if IS_PYDANTIC_V2 and isinstance(obj, pydantic.RootModel):
return self.encode(obj.root)
if isinstance(obj, pydantic.BaseModel):
items = self._iter_model_items(obj)
return {key: self.encode(value) for key, value in items}
if isinstance(obj, Mapping):
return {
key if isinstance(key, Enum) else str(key): self.encode(value)
for key, value in obj.items()
}
if isinstance(obj, Iterable):
return [self.encode(value) for value in obj]
raise ValueError(f"Cannot encode {obj!r}")
def _iter_model_items(
self, obj: pydantic.BaseModel
) -> Iterable[Tuple[str, Any]]:
exclude, keep_nulls = self.exclude, self.keep_nulls
get_model_field = get_model_fields(obj).get
for key, value in obj.__iter__():
field_info = get_model_field(key)
if field_info is not None:
key = field_info.alias or key
if key not in exclude and (value is not None or keep_nulls):
yield key, value
def _get_encoder(
obj: Any, custom_encoders: Mapping[type, SingleArgCallable]
) -> Optional[SingleArgCallable]:
encoder = custom_encoders.get(type(obj))
if encoder is not None:
return encoder
for cls, encoder in custom_encoders.items():
if isinstance(obj, cls):
return encoder
return None
|
