Functions inside models (#209)

samuelcolvin · web-flow · commit b33dc3e7649c · 2022-08-01T17:21:25.000+01:00
* Allow functions within model_class validators

* make fields_set flexible

* improve coverage

* fix error message test to support pypy
diff --git a/pydantic_core/_types.py b/pydantic_core/_types.py
@@ -112,7 +112,7 @@ class LiteralSchema(TypedDict):
 class ModelClassSchema(TypedDict):
     type: Literal['model-class']
     class_type: type
-    schema: TypedDictSchema
+    schema: Schema
     strict: NotRequired[bool]
     ref: NotRequired[str]
     config: NotRequired[Config]
diff --git a/src/validators/arguments.rs b/src/validators/arguments.rs
@@ -76,7 +76,7 @@ impl BuildValidator for ArgumentsValidator {
                 .get_as_req(intern!(py, "schema"))
                 .map_err(|err| SchemaError::new_err(format!("Argument \"{}\":\n  {}", name, err)))?;
 
-            let (validator, _) = build_validator(schema, config, build_context)?;
+            let validator = build_validator(schema, config, build_context)?;
 
             let default = arg.get_as(intern!(py, "default"))?;
             let default_factory = arg.get_as(intern!(py, "default_factory"))?;
@@ -102,11 +102,11 @@ impl BuildValidator for ArgumentsValidator {
             arguments,
             positional_args_count,
             var_args_validator: match schema.get_item(intern!(py, "var_args_schema")) {
-                Some(v) => Some(Box::new(build_validator(v, config, build_context)?.0)),
+                Some(v) => Some(Box::new(build_validator(v, config, build_context)?)),
                 None => None,
             },
             var_kwargs_validator: match schema.get_item(intern!(py, "var_kwargs_schema")) {
-                Some(v) => Some(Box::new(build_validator(v, config, build_context)?.0)),
+                Some(v) => Some(Box::new(build_validator(v, config, build_context)?)),
                 None => None,
             },
         }
diff --git a/src/validators/dict.rs b/src/validators/dict.rs
@@ -30,11 +30,11 @@ impl BuildValidator for DictValidator {
     ) -> PyResult<CombinedValidator> {
         let py = schema.py();
         let key_validator = match schema.get_item(intern!(py, "keys_schema")) {
-            Some(schema) => Box::new(build_validator(schema, config, build_context)?.0),
+            Some(schema) => Box::new(build_validator(schema, config, build_context)?),
             None => Box::new(AnyValidator::build(schema, config, build_context)?),
         };
         let value_validator = match schema.get_item(intern!(py, "values_schema")) {
-            Some(d) => Box::new(build_validator(d, config, build_context)?.0),
+            Some(d) => Box::new(build_validator(d, config, build_context)?),
             None => Box::new(AnyValidator::build(schema, config, build_context)?),
         };
         let name = format!(
diff --git a/src/validators/function.rs b/src/validators/function.rs
@@ -3,14 +3,13 @@ use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyAny, PyDict};
 
-use crate::build_tools::{py_error, SchemaDict};
+use crate::build_tools::SchemaDict;
 use crate::errors::{ErrorKind, PydanticValueError, ValError, ValResult, ValidationError};
 use crate::input::Input;
 use crate::recursion_guard::RecursionGuard;
 
 use super::{build_validator, BuildContext, BuildValidator, CombinedValidator, Extra, Validator};
 
-#[derive(Debug)]
 pub struct FunctionBuilder;
 
 impl BuildValidator for FunctionBuilder {
@@ -25,9 +24,9 @@ impl BuildValidator for FunctionBuilder {
         match mode {
             "before" => FunctionBeforeValidator::build(schema, config, build_context),
             "after" => FunctionAfterValidator::build(schema, config, build_context),
-            "plain" => FunctionPlainValidator::build(schema, config),
             "wrap" => FunctionWrapValidator::build(schema, config, build_context),
-            _ => py_error!("Unexpected function mode {:?}", mode),
+            // must be "plain"
+            _ => FunctionPlainValidator::build(schema, config),
         }
     }
 }
@@ -47,7 +46,7 @@ macro_rules! impl_build {
                 build_context: &mut BuildContext,
             ) -> PyResult<CombinedValidator> {
                 let py = schema.py();
-                let validator = build_validator(schema.get_as_req(intern!(py, "schema"))?, config, build_context)?.0;
+                let validator = build_validator(schema.get_as_req(intern!(py, "schema"))?, config, build_context)?;
                 let name = format!("{}[{}]", $name, validator.get_name());
                 Ok(Self {
                     validator: Box::new(validator),
@@ -110,6 +109,10 @@ impl Validator for FunctionBeforeValidator {
         &self.name
     }
 
+    fn ask(&self, question: &str) -> bool {
+        self.validator.ask(question)
+    }
+
     fn complete(&mut self, build_context: &BuildContext) -> PyResult<()> {
         self.validator.complete(build_context)
     }
@@ -143,6 +146,10 @@ impl Validator for FunctionAfterValidator {
         &self.name
     }
 
+    fn ask(&self, question: &str) -> bool {
+        self.validator.ask(question)
+    }
+
     fn complete(&mut self, build_context: &BuildContext) -> PyResult<()> {
         self.validator.complete(build_context)
     }
@@ -157,18 +164,14 @@ pub struct FunctionPlainValidator {
 impl FunctionPlainValidator {
     pub fn build(schema: &PyDict, config: Option<&PyDict>) -> PyResult<CombinedValidator> {
         let py = schema.py();
-        if schema.get_item(intern!(py, "schema")).is_some() {
-            py_error!("Plain functions should not include a sub-schema")
-        } else {
-            Ok(Self {
-                func: schema.get_as_req::<&PyAny>(intern!(py, "function"))?.into_py(py),
-                config: match config {
-                    Some(c) => c.into(),
-                    None => py.None(),
-                },
-            }
-            .into())
+        Ok(Self {
+            func: schema.get_as_req::<&PyAny>(intern!(py, "function"))?.into_py(py),
+            config: match config {
+                Some(c) => c.into(),
+                None => py.None(),
+            },
         }
+        .into())
     }
 }
 
@@ -236,6 +239,10 @@ impl Validator for FunctionWrapValidator {
         &self.name
     }
 
+    fn ask(&self, question: &str) -> bool {
+        self.validator.ask(question)
+    }
+
     fn complete(&mut self, build_context: &BuildContext) -> PyResult<()> {
         self.validator.complete(build_context)
     }
diff --git a/src/validators/list.rs b/src/validators/list.rs
@@ -28,7 +28,7 @@ macro_rules! generic_list_like_build {
         ) -> PyResult<CombinedValidator> {
             let py = schema.py();
             let item_validator = match schema.get_item(pyo3::intern!(py, "items_schema")) {
-                Some(d) => Some(Box::new(build_validator(d, config, build_context)?.0)),
+                Some(d) => Some(Box::new(build_validator(d, config, build_context)?)),
                 None => None,
             };
             let inner_name = item_validator.as_ref().map(|v| v.get_name()).unwrap_or("any");
diff --git a/src/validators/mod.rs b/src/validators/mod.rs
@@ -68,7 +68,7 @@ impl SchemaValidator {
         let schema = schema_obj.as_ref(py);
 
         let mut build_context = BuildContext::default();
-        let (mut validator, _) = build_validator(schema, config, &mut build_context)?;
+        let mut validator = build_validator(schema, config, &mut build_context)?;
         validator.complete(&build_context)?;
         let slots = build_context.into_slots()?;
         let title = validator.get_name().into_py(py);
@@ -211,7 +211,7 @@ impl SchemaValidator {
 
         let mut build_context = BuildContext::default();
         let validator = match build_validator(self_schema, None, &mut build_context) {
-            Ok((v, _)) => v,
+            Ok(v) => v,
             Err(err) => return Err(SchemaError::new_err(format!("Error building self-schema:\n  {}", err))),
         };
         Ok(Self {
@@ -255,7 +255,7 @@ fn build_single_validator<'a, T: BuildValidator>(
     schema_dict: &'a PyDict,
     config: Option<&'a PyDict>,
     build_context: &mut BuildContext,
-) -> PyResult<(CombinedValidator, &'a PyDict)> {
+) -> PyResult<CombinedValidator> {
     let py = schema_dict.py();
     let val: CombinedValidator = if let Some(schema_ref) = schema_dict.get_as::<String>(intern!(py, "ref"))? {
         let slot_id = build_context.prepare_slot(schema_ref)?;
@@ -269,7 +269,7 @@ fn build_single_validator<'a, T: BuildValidator>(
             .map_err(|err| SchemaError::new_err(format!("Error building \"{}\" validator:\n  {}", val_type, err)))?
     };
 
-    Ok((val, schema_dict))
+    Ok(val)
 }
 
 // macro to build the match statement for validator selection
@@ -290,7 +290,7 @@ pub fn build_validator<'a>(
     schema: &'a PyAny,
     config: Option<&'a PyDict>,
     build_context: &mut BuildContext,
-) -> PyResult<(CombinedValidator, &'a PyDict)> {
+) -> PyResult<CombinedValidator> {
     let py = schema.py();
     let dict: &PyDict = match schema.cast_as() {
         Ok(s) => s,
@@ -485,6 +485,14 @@ pub trait Validator: Send + Sync + Clone + Debug {
     /// this is used in the error location in unions, and in the top level message in `ValidationError`
     fn get_name(&self) -> &str;
 
+    /// allows validators to ask specific questions of sub-validators in a general way, could be extended
+    /// to do more, validators which don't know the question and have sub-validators
+    /// should return the result them in an `...iter().all(|v| v.ask(question))` way, ONLY
+    /// if they return the value of the sub-validator, e.g. functions, unions
+    fn ask(&self, _question: &str) -> bool {
+        false
+    }
+
     /// this method must be implemented for any validator which holds references to other validators,
     /// it is used by `RecursiveRefValidator` to set its name
     fn complete(&mut self, _build_context: &BuildContext) -> PyResult<()> {
diff --git a/src/validators/model_class.rs b/src/validators/model_class.rs
@@ -8,21 +8,21 @@ use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyTuple, PyType};
 use pyo3::{ffi, intern};
 
-use crate::build_tools::{py_error, SchemaDict};
+use crate::build_tools::SchemaDict;
 use crate::errors::{ErrorKind, ValError, ValResult};
 use crate::input::Input;
 use crate::recursion_guard::RecursionGuard;
 
-use super::typed_dict::TypedDictValidator;
 use super::{build_validator, BuildContext, BuildValidator, CombinedValidator, Extra, Validator};
 
 #[derive(Debug, Clone)]
 pub struct ModelClassValidator {
     strict: bool,
     revalidate: bool,
-    validator: TypedDictValidator,
+    validator: Box<CombinedValidator>,
     class: Py<PyType>,
     name: String,
+    expect_fields_set: bool,
 }
 
 impl BuildValidator for ModelClassValidator {
@@ -39,27 +39,21 @@ impl BuildValidator for ModelClassValidator {
 
         let class: &PyType = schema.get_as_req(intern!(py, "class_type"))?;
         let sub_schema: &PyAny = schema.get_as_req(intern!(py, "schema"))?;
-        let (comb_validator, td_schema) = build_validator(sub_schema, config, build_context)?;
+        let validator = build_validator(sub_schema, config, build_context)?;
 
-        if !td_schema.get_as(intern!(py, "return_fields_set"))?.unwrap_or(false) {
-            return py_error!("model-class inner schema should have 'return_fields_set' set to True");
-        }
-
-        let validator = match comb_validator {
-            CombinedValidator::TypedDict(tdv) => tdv,
-            _ => return py_error!("Wrong validator type, expected 'typed-dict' validator"),
-        };
+        let expect_fields_set = validator.ask("return_fields_set");
 
         Ok(Self {
             // we don't use is_strict here since we don't want validation to be strict in this case if
             // `config.strict` is set, only if this specific field is strict
             strict: schema.get_as(intern!(py, "strict"))?.unwrap_or(false),
             revalidate: config.get_as(intern!(py, "revalidate_models"))?.unwrap_or(false),
-            validator,
+            validator: Box::new(validator),
             class: class.into(),
             // Get the class's `__name__`, not using `class.name()` since it uses `__qualname__`
             // which is not what we want here
             name: class.getattr(intern!(py, "__name__"))?.extract()?,
+            expect_fields_set,
         }
         .into())
     }
@@ -79,9 +73,13 @@ impl Validator for ModelClassValidator {
             if self.revalidate {
                 let fields_set = input.get_attr(intern!(py, "__fields_set__"));
                 let output = self.validator.validate(py, input, extra, slots, recursion_guard)?;
-                let (model_dict, validation_fields_set): (&PyAny, &PyAny) = output.extract(py)?;
-                let fields_set = fields_set.unwrap_or(validation_fields_set);
-                Ok(self.create_class(py, model_dict, fields_set)?)
+                if self.expect_fields_set {
+                    let (model_dict, validation_fields_set): (&PyAny, &PyAny) = output.extract(py)?;
+                    let fields_set = fields_set.unwrap_or(validation_fields_set);
+                    Ok(self.create_class(py, model_dict, Some(fields_set))?)
+                } else {
+                    Ok(self.create_class(py, output.as_ref(py), fields_set)?)
+                }
             } else {
                 Ok(input.to_object(py))
             }
@@ -94,8 +92,12 @@ impl Validator for ModelClassValidator {
             ))
         } else {
             let output = self.validator.validate(py, input, extra, slots, recursion_guard)?;
-            let (model_dict, fields_set): (&PyAny, &PyAny) = output.extract(py)?;
-            Ok(self.create_class(py, model_dict, fields_set)?)
+            if self.expect_fields_set {
+                let (model_dict, fields_set): (&PyAny, &PyAny) = output.extract(py)?;
+                Ok(self.create_class(py, model_dict, Some(fields_set))?)
+            } else {
+                Ok(self.create_class(py, output.as_ref(py), None)?)
+            }
         }
     }
 
@@ -105,7 +107,7 @@ impl Validator for ModelClassValidator {
 }
 
 impl ModelClassValidator {
-    fn create_class(&self, py: Python, model_dict: &PyAny, fields_set: &PyAny) -> PyResult<PyObject> {
+    fn create_class(&self, py: Python, model_dict: &PyAny, fields_set: Option<&PyAny>) -> PyResult<PyObject> {
         // based on the following but with the second argument of new_func set to an empty tuple as required
         // https://github.com/PyO3/pyo3/blob/d2caa056e9aacc46374139ef491d112cb8af1a25/src/pyclass_init.rs#L35-L77
         let args = PyTuple::empty(py);
@@ -126,7 +128,9 @@ impl ModelClassValidator {
 
         let instance_ref = instance.as_ref(py);
         force_setattr(py, instance_ref, intern!(py, "__dict__"), model_dict)?;
-        force_setattr(py, instance_ref, intern!(py, "__fields_set__"), fields_set)?;
+        if let Some(fields_set) = fields_set {
+            force_setattr(py, instance_ref, intern!(py, "__fields_set__"), fields_set)?;
+        }
 
         Ok(instance)
     }
diff --git a/src/validators/nullable.rs b/src/validators/nullable.rs
@@ -24,7 +24,7 @@ impl BuildValidator for NullableValidator {
         build_context: &mut BuildContext,
     ) -> PyResult<CombinedValidator> {
         let schema: &PyAny = schema.get_as_req(intern!(schema.py(), "schema"))?;
-        let validator = Box::new(build_validator(schema, config, build_context)?.0);
+        let validator = Box::new(build_validator(schema, config, build_context)?);
         let name = format!("{}[{}]", Self::EXPECTED_TYPE, validator.get_name());
         Ok(Self { validator, name }.into())
     }
diff --git a/src/validators/tuple.rs b/src/validators/tuple.rs
@@ -93,15 +93,15 @@ impl TuplePositionalValidator {
         let items: &PyList = schema.get_as_req(intern!(py, "items_schema"))?;
         let validators: Vec<CombinedValidator> = items
             .iter()
-            .map(|item| build_validator(item, config, build_context).map(|result| result.0))
+            .map(|item| build_validator(item, config, build_context))
             .collect::<PyResult<Vec<CombinedValidator>>>()?;
 
         let descr = validators.iter().map(|v| v.get_name()).collect::<Vec<_>>().join(", ");
         Ok(Self {
             strict: is_strict(schema, config)?,
             items_validators: validators,
             extra_validator: match schema.get_item(intern!(py, "extra_schema")) {
-                Some(v) => Some(Box::new(build_validator(v, config, build_context)?.0)),
+                Some(v) => Some(Box::new(build_validator(v, config, build_context)?)),
                 None => None,
             },
             name: format!("tuple[{}]", descr),
diff --git a/src/validators/typed_dict.rs b/src/validators/typed_dict.rs
@@ -94,7 +94,7 @@ impl BuildValidator for TypedDictValidator {
         let extra_validator = match schema.get_item(intern!(py, "extra_validator")) {
             Some(v) => {
                 if check_extra && !forbid_extra {
-                    Some(Box::new(build_validator(v, config, build_context)?.0))
+                    Some(Box::new(build_validator(v, config, build_context)?))
                 } else {
                     return py_error!("extra_validator can only be used if extra_behavior=allow");
                 }
@@ -175,7 +175,7 @@ impl BuildValidator for TypedDictValidator {
                 lookup_key,
                 name_pystring: PyString::intern(py, field_name).into(),
                 validator: match build_validator(schema, config, build_context) {
-                    Ok((v, _)) => v,
+                    Ok(v) => v,
                     Err(err) => return py_error!("Field \"{}\":\n  {}", field_name, err),
                 },
                 required,
@@ -376,6 +376,14 @@ impl Validator for TypedDictValidator {
         Self::EXPECTED_TYPE
     }
 
+    fn ask(&self, question: &str) -> bool {
+        if question == "return_fields_set" {
+            self.return_fields_set
+        } else {
+            false
+        }
+    }
+
     fn complete(&mut self, build_context: &BuildContext) -> PyResult<()> {
         self.fields
             .iter_mut()
diff --git a/src/validators/union.rs b/src/validators/union.rs
diff --git a/tests/benchmarks/test_micro_benchmarks.py b/tests/benchmarks/test_micro_benchmarks.py
diff --git a/tests/validators/test_model_class.py b/tests/validators/test_model_class.py

Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ impl BuildValidator for NullableValidator {`
`24`	`24`	`build_context: &mut BuildContext,`
`25`	`25`	`) -> PyResult<CombinedValidator> {`
`26`	`26`	`let schema: &PyAny = schema.get_as_req(intern!(schema.py(), "schema"))?;`
`27`		`- let validator = Box::new(build_validator(schema, config, build_context)?.0);`
	`27`	`+ let validator = Box::new(build_validator(schema, config, build_context)?);`
`28`	`28`	`let name = format!("{}[{}]", Self::EXPECTED_TYPE, validator.get_name());`
`29`	`29`	`Ok(Self { validator, name }.into())`
`30`	`30`	`}`