Chain type (#255)

* implement chain, chain2 and chain3

* remove chain2 and chain3, more tests

* benchmarks and test flatten

* fix schema

* improve test coverage

Author: samuelcolvin

pydantic_core/_pydantic_core.pyi

@@ -12,6 +12,7 @@ __all__ = '__version__', 'SchemaValidator', 'SchemaError', 'ValidationError', 'P
 __version__: str
 
 class SchemaValidator:
+    title: str
     def __init__(self, schema: CoreSchema, config: 'CoreConfig | None' = None) -> None: ...
     def validate_python(self, input: Any, strict: 'bool | None' = None, context: Any = None) -> Any: ...
     def isinstance_python(self, input: Any, strict: 'bool | None' = None, context: Any = None) -> bool: ...

pydantic_core/core_schema.py

@@ -315,6 +315,12 @@ class WithDefaultSchema(TypedDict, total=False):
     ref: str
 
 
+class ChainSchema(TypedDict):
+    type: Literal['chain']
+    steps: List[CoreSchema]
+    ref: NotRequired[str]
+
+
 CoreSchema = Union[
     AnySchema,
     BoolSchema,
@@ -347,4 +353,5 @@ class WithDefaultSchema(TypedDict, total=False):
     ArgumentsSchema,
     CallSchema,
     WithDefaultSchema,
+    ChainSchema,
 ]

src/validators/chain.rs

@@ -0,0 +1,101 @@
+use pyo3::intern;
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyList};
+
+use crate::build_tools::{py_error, SchemaDict};
+use crate::errors::ValResult;
+use crate::input::Input;
+use crate::questions::Question;
+use crate::recursion_guard::RecursionGuard;
+
+use super::{build_validator, BuildContext, BuildValidator, CombinedValidator, Extra, Validator};
+
+#[derive(Debug, Clone)]
+pub struct ChainValidator {
+    steps: Vec<CombinedValidator>,
+    name: String,
+}
+
+impl BuildValidator for ChainValidator {
+    const EXPECTED_TYPE: &'static str = "chain";
+
+    fn build(
+        schema: &PyDict,
+        config: Option<&PyDict>,
+        build_context: &mut BuildContext,
+    ) -> PyResult<CombinedValidator> {
+        let steps: Vec<CombinedValidator> = schema
+            .get_as_req::<&PyList>(intern!(schema.py(), "steps"))?
+            .iter()
+            .map(|step| build_validator_steps(step, config, build_context))
+            .collect::<PyResult<Vec<Vec<CombinedValidator>>>>()?
+            .into_iter()
+            .flatten()
+            .collect::<Vec<CombinedValidator>>();
+
+        match steps.len() {
+            0 => py_error!("One or more steps are required for a chain validator"),
+            1 => {
+                let step = steps.into_iter().next().unwrap();
+                Ok(step)
+            }
+            _ => {
+                let descr = steps.iter().map(|v| v.get_name()).collect::<Vec<_>>().join(",");
+
+                Ok(Self {
+                    steps,
+                    name: format!("{}[{}]", Self::EXPECTED_TYPE, descr),
+                }
+                .into())
+            }
+        }
+    }
+}
+
+// either a vec of the steps from a nested `ChainValidator`, or a length-1 vec containing the validator
+// to be flattened into `steps` above
+fn build_validator_steps<'a>(
+    step: &'a PyAny,
+    config: Option<&'a PyDict>,
+    build_context: &mut BuildContext,
+) -> PyResult<Vec<CombinedValidator>> {
+    let validator = build_validator(step, config, build_context)?;
+    if let CombinedValidator::Chain(chain_validator) = validator {
+        Ok(chain_validator.steps)
+    } else {
+        Ok(vec![validator])
+    }
+}
+
+impl Validator for ChainValidator {
+    fn validate<'s, 'data>(
+        &'s self,
+        py: Python<'data>,
+        input: &'data impl Input<'data>,
+        extra: &Extra,
+        slots: &'data [CombinedValidator],
+        recursion_guard: &'s mut RecursionGuard,
+    ) -> ValResult<'data, PyObject> {
+        let mut steps_iter = self.steps.iter();
+        let first_step = steps_iter.next().unwrap();
+        let value = first_step.validate(py, input, extra, slots, recursion_guard)?;
+
+        steps_iter.try_fold(value, |v, step| {
+            step.validate(py, v.into_ref(py), extra, slots, recursion_guard)
+        })
+    }
+
+    fn get_name(&self) -> &str {
+        &self.name
+    }
+
+    fn ask(&self, question: &Question) -> bool {
+        // any makes more sense since at the moment we only use ask for "return_fields_set", might need
+        // more complex logic in future
+        self.steps.iter().any(|v| v.ask(question))
+    }
+
+    fn complete(&mut self, build_context: &BuildContext) -> PyResult<()> {
+        self.steps.iter_mut().try_for_each(|v| v.complete(build_context))
+    }
+}

src/validators/mod.rs

@@ -22,6 +22,7 @@ mod bool;
 mod bytes;
 mod call;
 mod callable;
+mod chain;
 mod date;
 mod datetime;
 mod dict;
@@ -51,6 +52,7 @@ pub struct SchemaValidator {
     validator: CombinedValidator,
     slots: Vec<CombinedValidator>,
     schema: PyObject,
+    #[pyo3(get)]
     title: PyObject,
 }
 
@@ -380,6 +382,8 @@ pub fn build_validator<'a>(
         arguments::ArgumentsValidator,
         // default value
         with_default::WithDefaultValidator,
+        // chain validators
+        chain::ChainValidator,
     )
 }
 
@@ -490,6 +494,8 @@ pub enum CombinedValidator {
     Arguments(arguments::ArgumentsValidator),
     // default value
     WithDefault(with_default::WithDefaultValidator),
+    // chain validators
+    Chain(chain::ChainValidator),
 }
 
 /// This trait must be implemented by all validators, it allows various validators to be accessed consistently,

tests/benchmarks/test_micro_benchmarks.py

@@ -5,6 +5,7 @@
 import os
 import platform
 from datetime import date, datetime, timedelta, timezone
+from decimal import Decimal
 from typing import Dict, FrozenSet, List, Optional, Set, Union
 
 import pytest
@@ -890,3 +891,66 @@ def test_with_default(benchmark):
     def t():
         v.validate_python({'name': 'Foo'})
         v.validate_python({})
+
+
+@pytest.mark.benchmark(group='chain')
+def test_chain_list(benchmark):
+    validator = SchemaValidator(
+        {
+            'type': 'chain',
+            'steps': [
+                {'type': 'str'},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: Decimal(v)},
+            ],
+        }
+    )
+    assert validator.validate_python('42.42') == Decimal('42.42')
+
+    benchmark(validator.validate_python, '42.42')
+
+
+@pytest.mark.benchmark(group='chain')
+def test_chain_function(benchmark):
+    validator = SchemaValidator(
+        {'type': 'function', 'mode': 'after', 'schema': {'type': 'str'}, 'function': lambda v, **kwargs: Decimal(v)}
+    )
+    assert validator.validate_python('42.42') == Decimal('42.42')
+
+    benchmark(validator.validate_python, '42.42')
+
+
+@pytest.mark.benchmark(group='chain-functions')
+def test_chain_two_functions(benchmark):
+    validator = SchemaValidator(
+        {
+            'type': 'chain',
+            'steps': [
+                {'type': 'str'},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: Decimal(v)},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: v * 2},
+            ],
+        }
+    )
+    assert validator.validate_python('42.42') == Decimal('84.84')
+
+    benchmark(validator.validate_python, '42.42')
+
+
+@pytest.mark.benchmark(group='chain-functions')
+def test_chain_nested_functions(benchmark):
+    validator = SchemaValidator(
+        {
+            'type': 'function',
+            'schema': {
+                'type': 'function',
+                'schema': {'type': 'str'},
+                'mode': 'after',
+                'function': lambda v, **kwargs: Decimal(v),
+            },
+            'mode': 'after',
+            'function': lambda v, **kwargs: v * 2,
+        }
+    )
+    assert validator.validate_python('42.42') == Decimal('84.84')
+
+    benchmark(validator.validate_python, '42.42')

tests/validators/test_chain.py

@@ -0,0 +1,146 @@
+import re
+from decimal import Decimal
+
+import pytest
+
+from pydantic_core import SchemaError, SchemaValidator, ValidationError
+
+from ..conftest import PyAndJson, plain_repr
+
+
+def test_chain():
+    validator = SchemaValidator(
+        {
+            'type': 'chain',
+            'steps': [
+                {'type': 'str'},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: Decimal(v)},
+            ],
+        }
+    )
+
+    assert validator.validate_python('1.44') == Decimal('1.44')
+    assert validator.validate_python(b'1.44') == Decimal('1.44')
+
+
+def test_chain_many():
+    validator = SchemaValidator(
+        {
+            'type': 'chain',
+            'steps': [
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-1'},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-2'},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-3'},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-4'},
+            ],
+        }
+    )
+
+    assert validator.validate_python('input') == 'input-1-2-3-4'
+
+
+def test_chain_error():
+    validator = SchemaValidator({'type': 'chain', 'steps': [{'type': 'str'}, {'type': 'int'}]})
+
+    assert validator.validate_python('123') == 123
+    assert validator.validate_python(b'123') == 123
+
+    with pytest.raises(ValidationError) as exc_info:
+        validator.validate_python('abc')
+    # insert_assert(exc_info.value.errors())
+    assert exc_info.value.errors() == [
+        {
+            'kind': 'int_parsing',
+            'loc': [],
+            'message': 'Input should be a valid integer, unable to parse string as an integer',
+            'input_value': 'abc',
+        }
+    ]
+
+
+@pytest.mark.parametrize(
+    'input_value,expected', [('1.44', Decimal('1.44')), (1, Decimal(1)), (1.44, pytest.approx(1.44))]
+)
+def test_json(py_and_json: PyAndJson, input_value, expected):
+    validator = py_and_json(
+        {
+            'type': 'chain',
+            'steps': [
+                {'type': 'union', 'choices': [{'type': 'str'}, {'type': 'float'}]},
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: Decimal(v)},
+            ],
+        }
+    )
+    output = validator.validate_test(input_value)
+    assert output == expected
+    assert isinstance(output, Decimal)
+
+
+def test_flatten():
+    validator = SchemaValidator(
+        {
+            'type': 'chain',
+            'steps': [
+                {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-1'},
+                {
+                    'type': 'chain',
+                    'steps': [
+                        {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-2'},
+                        {'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-3'},
+                    ],
+                },
+            ],
+        }
+    )
+
+    assert validator.validate_python('input') == 'input-1-2-3'
+    assert validator.title == 'chain[function-plain,function-plain,function-plain]'
+
+
+def test_chain_empty():
+    with pytest.raises(SchemaError, match='One or more steps are required for a chain validator'):
+        SchemaValidator({'type': 'chain', 'steps': []})
+
+
+def test_chain_one():
+    validator = SchemaValidator(
+        {'type': 'chain', 'steps': [{'type': 'function', 'mode': 'plain', 'function': lambda v, **kwargs: f'{v}-1'}]}
+    )
+    assert validator.validate_python('input') == 'input-1'
+    assert validator.title == 'function-plain'
+
+
+def test_ask():
+    class MyModel:
+        __slots__ = '__dict__', '__fields_set__'
+
+    calls = []
+
+    def f(input_value, **kwargs):
+        calls.append(input_value)
+        return input_value
+
+    v = SchemaValidator(
+        {
+            'type': 'new-class',
+            'class_type': MyModel,
+            'schema': {
+                'type': 'chain',
+                'steps': [
+                    {
+                        'type': 'typed-dict',
+                        'return_fields_set': True,
+                        'fields': {'field_a': {'schema': {'type': 'str'}}},
+                    },
+                    {'type': 'function', 'mode': 'plain', 'function': f},
+                ],
+            },
+        }
+    )
+    assert re.search('expect_fields_set:(true|false)', plain_repr(v)).group(1) == 'true'
+    m = v.validate_python({'field_a': 'abc'})
+    assert isinstance(m, MyModel)
+    assert m.field_a == 'abc'
+    assert m.__fields_set__ == {'field_a'}
+    # insert_assert(calls)
+    assert calls == [({'field_a': 'abc'}, {'field_a'})]