Skip to content

Commit 193b998

Browse files
samuelcolvinsamuelcolvin
committed
customising fields
1 parent ce8d618 commit 193b998

File tree

9 files changed

+123
-115
lines changed

9 files changed

+123
-115
lines changed

‎src/serializers/fields.rs

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
use std::borrow::Cow;
2+
3+
use pyo3::prelude::*;
4+
use pyo3::types::{PyDict, PyString};
5+
6+
use super::extra::Extra;
7+
use super::shared::{CombinedSerializer, TypeSerializer};
8+
9+
/// representation of a field for serialization, used by `TypedDictSerializer` and `ModelFieldsSerializer`,
10+
/// and maybe more
11+
#[derive(Debug, Clone)]
12+
pub(super) struct FieldSerializer {
13+
pub key_py: Py<PyString>,
14+
pub alias: Option<String>,
15+
pub alias_py: Option<Py<PyString>>,
16+
// None serializer means exclude
17+
pub serializer: Option<CombinedSerializer>,
18+
pub required: bool,
19+
}
20+
21+
impl FieldSerializer {
22+
pub fn new(
23+
py: Python,
24+
key_py: Py<PyString>,
25+
alias: Option<String>,
26+
serializer: Option<CombinedSerializer>,
27+
required: bool,
28+
) -> Self {
29+
let alias_py = alias.as_ref().map(|alias| PyString::new(py, alias.as_str()).into());
30+
Self {
31+
key_py,
32+
alias,
33+
alias_py,
34+
serializer,
35+
required,
36+
}
37+
}
38+
39+
pub fn get_key_py<'py>(&'py self, py: Python<'py>, extra: &Extra) -> &'py PyAny {
40+
if extra.by_alias {
41+
if let Some(ref alias_py) = self.alias_py {
42+
return alias_py.as_ref(py);
43+
}
44+
}
45+
self.key_py.as_ref(py)
46+
}
47+
48+
pub fn get_key_json<'a>(&'a self, key_str: &'a str, extra: &Extra) -> Cow<'a, str> {
49+
if extra.by_alias {
50+
if let Some(ref alias) = self.alias {
51+
return Cow::Borrowed(alias.as_str());
52+
}
53+
}
54+
Cow::Borrowed(key_str)
55+
}
56+
57+
pub fn to_python(
58+
&self,
59+
output_dict: &PyDict,
60+
value: &PyAny,
61+
next_include: Option<&PyAny>,
62+
next_exclude: Option<&PyAny>,
63+
extra: &Extra,
64+
) -> PyResult<()> {
65+
if let Some(ref serializer) = self.serializer {
66+
if !exclude_default(value, extra, serializer)? {
67+
let value = serializer.to_python(value, next_include, next_exclude, extra)?;
68+
let output_key = self.get_key_py(output_dict.py(), extra);
69+
output_dict.set_item(output_key, value)?;
70+
}
71+
}
72+
Ok(())
73+
}
74+
}
75+
76+
pub(super) fn exclude_default(value: &PyAny, extra: &Extra, serializer: &CombinedSerializer) -> PyResult<bool> {
77+
if extra.exclude_defaults {
78+
if let Some(default) = serializer.get_default(value.py())? {
79+
if value.eq(default)? {
80+
return Ok(true);
81+
}
82+
}
83+
}
84+
Ok(false)
85+
}

‎src/serializers/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ mod computed_fields;
1818
mod config;
1919
mod errors;
2020
mod extra;
21+
mod fields;
2122
mod filter;
2223
mod infer;
2324
mod ob_type;

‎src/serializers/type_serializers/dataclass.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ use crate::build_tools::{py_error_type, SchemaDict};
88
use crate::definitions::DefinitionsBuilder;
99

1010
use super::model::ModelSerializer;
11-
use super::typed_dict::{FieldSerializer, TypedDictSerializer};
12-
use super::{BuildSerializer, CombinedSerializer, ComputedFields};
11+
use super::typed_dict::TypedDictSerializer;
12+
use super::{BuildSerializer, CombinedSerializer, ComputedFields, FieldSerializer};
1313

1414
pub struct DataclassArgsBuilder;
1515

‎src/serializers/type_serializers/function.rs

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,14 @@ use serde::ser::Error;
1111

1212
use crate::build_tools::{function_name, py_error_type, SchemaDict};
1313
use crate::definitions::DefinitionsBuilder;
14-
use crate::serializers::extra::{ExtraOwned, SerMode};
15-
use crate::serializers::filter::AnyFilter;
1614
use crate::{PydanticOmit, PydanticSerializationUnexpectedValue};
1715

1816
use super::format::WhenUsed;
1917

2018
use super::{
2119
infer_json_key, infer_json_key_known, infer_serialize, infer_serialize_known, infer_to_python,
22-
infer_to_python_known, py_err_se_err, BuildSerializer, CombinedSerializer, Extra, ObType,
23-
PydanticSerializationError, TypeSerializer,
20+
infer_to_python_known, py_err_se_err, AnyFilter, BuildSerializer, CombinedSerializer, Extra, ExtraOwned, ObType,
21+
PydanticSerializationError, SerMode, TypeSerializer,
2422
};
2523

2624
pub struct FunctionBeforeSerializerBuilder;

‎src/serializers/type_serializers/mod.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,9 @@ pub mod with_default;
2626
pub(self) use super::computed_fields::ComputedFields;
2727
pub(self) use super::config::utf8_py_error;
2828
pub(self) use super::errors::{py_err_se_err, PydanticSerializationError};
29-
pub(self) use super::extra::{Extra, ExtraOwned, SerMode};
30-
pub(self) use super::filter::SchemaFilter;
29+
pub(self) use super::extra::{Extra, ExtraOwned, SerCheck, SerMode};
30+
pub(self) use super::fields::{exclude_default, FieldSerializer};
31+
pub(self) use super::filter::{AnyFilter, SchemaFilter};
3132
pub(self) use super::infer::{
3233
infer_json_key, infer_json_key_known, infer_serialize, infer_serialize_known, infer_to_python,
3334
infer_to_python_known, SerializeInfer,

‎src/serializers/type_serializers/model.rs

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,11 @@ use ahash::AHashMap;
88

99
use crate::build_tools::{py_error_type, ExtraBehavior, SchemaDict};
1010
use crate::definitions::DefinitionsBuilder;
11-
use crate::serializers::computed_fields::ComputedFields;
12-
use crate::serializers::extra::SerCheck;
13-
use crate::serializers::infer::{infer_serialize, infer_to_python};
14-
use crate::serializers::ob_type::ObType;
15-
use crate::serializers::type_serializers::typed_dict::{FieldSerializer, TypedDictSerializer};
1611

12+
use super::typed_dict::TypedDictSerializer;
1713
use super::{
18-
infer_json_key, infer_json_key_known, object_to_dict, py_err_se_err, BuildSerializer, CombinedSerializer, Extra,
19-
TypeSerializer,
14+
infer_json_key, infer_json_key_known, infer_serialize, infer_to_python, object_to_dict, py_err_se_err,
15+
BuildSerializer, CombinedSerializer, ComputedFields, Extra, FieldSerializer, ObType, SerCheck, TypeSerializer,
2016
};
2117

2218
pub struct ModelFieldsBuilder;

‎src/serializers/type_serializers/other.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,9 @@ use pyo3::types::{PyDict, PyList};
44

55
use crate::build_tools::{py_err, SchemaDict};
66
use crate::definitions::DefinitionsBuilder;
7-
use crate::serializers::shared::CombinedSerializer;
87

98
use super::any::AnySerializer;
10-
use super::BuildSerializer;
9+
use super::{BuildSerializer, CombinedSerializer};
1110

1211
pub struct ChainBuilder;
1312

‎src/serializers/type_serializers/typed_dict.rs

Lines changed: 25 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -4,73 +4,27 @@ use pyo3::intern;
44
use pyo3::prelude::*;
55
use pyo3::types::{PyDict, PyString};
66

7-
use ahash::{AHashMap, AHashSet};
7+
use ahash::AHashMap;
88
use serde::ser::SerializeMap;
99

1010
use crate::build_tools::{py_error_type, schema_or_config, ExtraBehavior, SchemaDict};
1111
use crate::definitions::DefinitionsBuilder;
1212
use crate::PydanticSerializationUnexpectedValue;
1313

1414
use super::{
15-
infer_json_key, infer_serialize, infer_to_python, py_err_se_err, BuildSerializer, CombinedSerializer,
16-
ComputedFields, Extra, PydanticSerializer, SchemaFilter, SerializeInfer, TypeSerializer,
15+
exclude_default, infer_json_key, infer_serialize, infer_to_python, py_err_se_err, BuildSerializer,
16+
CombinedSerializer, ComputedFields, Extra, FieldSerializer, PydanticSerializer, SchemaFilter, SerializeInfer,
17+
TypeSerializer,
1718
};
1819

19-
/// representation of a field for serialization, used by `TypedDictSerializer` and `ModelFieldsSerializer`
20-
#[derive(Debug, Clone)]
21-
pub(super) struct FieldSerializer {
22-
key_py: Py<PyString>,
23-
alias: Option<String>,
24-
alias_py: Option<Py<PyString>>,
25-
// None serializer means exclude
26-
serializer: Option<CombinedSerializer>,
27-
required: bool,
28-
}
29-
30-
impl FieldSerializer {
31-
pub(super) fn new(
32-
py: Python,
33-
key_py: Py<PyString>,
34-
alias: Option<String>,
35-
serializer: Option<CombinedSerializer>,
36-
required: bool,
37-
) -> Self {
38-
let alias_py = alias.as_ref().map(|alias| PyString::new(py, alias.as_str()).into());
39-
Self {
40-
key_py,
41-
alias,
42-
alias_py,
43-
serializer,
44-
required,
45-
}
46-
}
47-
48-
fn get_key_py<'py>(&'py self, py: Python<'py>, extra: &Extra) -> &'py PyAny {
49-
if extra.by_alias {
50-
if let Some(ref alias_py) = self.alias_py {
51-
return alias_py.as_ref(py);
52-
}
53-
}
54-
self.key_py.as_ref(py)
55-
}
56-
57-
fn get_key_json<'a>(&'a self, key_str: &'a str, extra: &Extra) -> Cow<'a, str> {
58-
if extra.by_alias {
59-
if let Some(ref alias) = self.alias {
60-
return Cow::Borrowed(alias.as_str());
61-
}
62-
}
63-
Cow::Borrowed(key_str)
64-
}
65-
}
66-
6720
#[derive(Debug, Clone)]
6821
pub struct TypedDictSerializer {
6922
fields: AHashMap<String, FieldSerializer>,
7023
computed_fields: Option<ComputedFields>,
7124
include_extra: bool,
7225
// isize because we look up filter via `.hash()` which returns an isize
7326
filter: SchemaFilter<isize>,
27+
required_fields: usize,
7428
}
7529

7630
impl BuildSerializer for TypedDictSerializer {
@@ -126,11 +80,13 @@ impl TypedDictSerializer {
12680
include_extra: bool,
12781
computed_fields: Option<ComputedFields>,
12882
) -> Self {
83+
let required_fields = fields.values().filter(|f| f.required).count();
12984
Self {
13085
fields,
13186
include_extra,
13287
filter: SchemaFilter::default(),
13388
computed_fields,
89+
required_fields,
13490
}
13591
}
13692

@@ -140,17 +96,6 @@ impl TypedDictSerializer {
14096
Some(ref computed_fields) => computed_fields.len(),
14197
}
14298
}
143-
144-
fn exclude_default(&self, value: &PyAny, extra: &Extra, serializer: &CombinedSerializer) -> PyResult<bool> {
145-
if extra.exclude_defaults {
146-
if let Some(default) = serializer.get_default(value.py())? {
147-
if value.eq(default)? {
148-
return Ok(true);
149-
}
150-
}
151-
}
152-
Ok(false)
153-
}
15499
}
155100

156101
impl TypeSerializer for TypedDictSerializer {
@@ -173,11 +118,7 @@ impl TypeSerializer for TypedDictSerializer {
173118
Ok(py_dict) => {
174119
// NOTE! we maintain the order of the input dict assuming that's right
175120
let output_dict = PyDict::new(py);
176-
let mut used_fields = if td_extra.check.enabled() {
177-
Some(AHashSet::with_capacity(self.fields.len()))
178-
} else {
179-
None
180-
};
121+
let mut used_req_fields: usize = 0;
181122

182123
for (key, value) in py_dict {
183124
if extra.exclude_none && value.is_none() {
@@ -191,19 +132,10 @@ impl TypeSerializer for TypedDictSerializer {
191132
if let Ok(key_py_str) = key.downcast::<PyString>() {
192133
let key_str = key_py_str.to_str()?;
193134
if let Some(field) = self.fields.get(key_str) {
194-
let serializer = match field.serializer {
195-
Some(ref serializer) => serializer,
196-
None => continue,
197-
};
198-
if self.exclude_default(value, &extra, serializer)? {
199-
continue;
200-
}
201-
let value = serializer.to_python(value, next_include, next_exclude, &extra)?;
202-
let output_key = field.get_key_py(py, &extra);
203-
output_dict.set_item(output_key, value)?;
135+
field.to_python(output_dict, value, next_include, next_exclude, &extra)?;
204136

205-
if let Some(ref mut used_fields) = used_fields {
206-
used_fields.insert(key_str);
137+
if field.required {
138+
used_req_fields += 1;
207139
}
208140
continue;
209141
}
@@ -217,14 +149,8 @@ impl TypeSerializer for TypedDictSerializer {
217149
}
218150
}
219151
}
220-
if let Some(ref used_fields) = used_fields {
221-
let unused_fields = self
222-
.fields
223-
.iter()
224-
.any(|(k, v)| v.required && !used_fields.contains(k.as_str()));
225-
if unused_fields {
226-
return Err(PydanticSerializationUnexpectedValue::new_err(None));
227-
}
152+
if td_extra.check.enabled() && self.required_fields != used_req_fields {
153+
return Err(PydanticSerializationUnexpectedValue::new_err(None));
228154
}
229155
if let Some(ref computed_fields) = self.computed_fields {
230156
if let Some(model) = td_extra.model {
@@ -283,16 +209,19 @@ impl TypeSerializer for TypedDictSerializer {
283209
if let Ok(key_py_str) = key.downcast::<PyString>() {
284210
let key_str = key_py_str.to_str().map_err(py_err_se_err)?;
285211
if let Some(field) = self.fields.get(key_str) {
286-
let serializer = match field.serializer {
287-
Some(ref serializer) => serializer,
288-
None => continue,
289-
};
290-
if self.exclude_default(value, &extra, serializer).map_err(py_err_se_err)? {
291-
continue;
212+
if let Some(ref serializer) = field.serializer {
213+
if !exclude_default(value, &extra, serializer).map_err(py_err_se_err)? {
214+
let s = PydanticSerializer::new(
215+
value,
216+
serializer,
217+
next_include,
218+
next_exclude,
219+
&extra,
220+
);
221+
let output_key = field.get_key_json(key_str, &extra);
222+
map.serialize_entry(&output_key, &s)?;
223+
}
292224
}
293-
let output_key = field.get_key_json(key_str, &extra);
294-
let s = PydanticSerializer::new(value, serializer, next_include, next_exclude, &extra);
295-
map.serialize_entry(&output_key, &s)?;
296225
continue;
297226
}
298227
}

‎src/serializers/type_serializers/union.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,11 @@ use std::borrow::Cow;
55

66
use crate::build_tools::{py_err, SchemaDict};
77
use crate::definitions::DefinitionsBuilder;
8-
use crate::serializers::extra::SerCheck;
98
use crate::PydanticSerializationUnexpectedValue;
109

1110
use super::{
1211
infer_json_key, infer_serialize, infer_to_python, py_err_se_err, BuildSerializer, CombinedSerializer, Extra,
13-
TypeSerializer,
12+
SerCheck, TypeSerializer,
1413
};
1514

1615
#[derive(Debug, Clone)]

0 commit comments

Comments
 (0)