Add IsSubclassValidator (#301)

* add IsSubclassValidator

* tweak type hint

Author: samuelcolvin

pydantic_core/core_schema.py

@@ -378,6 +378,20 @@ def is_instance_schema(
     )
 
 
+class IsSubclassSchema(TypedDict, total=False):
+    type: Required[Literal['is-subclass']]
+    cls: Required[Type[Any]]
+    cls_repr: str
+    ref: str
+    extra: Any
+
+
+def is_subclass_schema(
+    cls: Type[Any], *, cls_repr: str | None = None, ref: str | None = None, extra: Any = None
+) -> IsInstanceSchema:
+    return dict_not_none(type='is-subclass', cls=cls, cls_repr=cls_repr, ref=ref, extra=extra)
+
+
 class CallableSchema(TypedDict, total=False):
     type: Required[Literal['callable']]
     ref: str
@@ -1027,6 +1041,7 @@ def json_schema(schema: CoreSchema | None = None, *, ref: str | None = None, ext
     TimedeltaSchema,
     LiteralSchema,
     IsInstanceSchema,
+    IsSubclassSchema,
     CallableSchema,
     ListSchema,
     TuplePositionalSchema,
@@ -1118,6 +1133,7 @@ def json_schema(schema: CoreSchema | None = None, *, ref: str | None = None, ext
     'time_delta_parsing',
     'frozen_set_type',
     'is_instance_of',
+    'is_subclass_of',
     'callable_type',
     'union_tag_invalid',
     'union_tag_not_found',

src/errors/kinds.rs

@@ -279,6 +279,10 @@ pub enum ErrorKind {
     IsInstanceOf {
         class: String,
     },
+    #[strum(message = "Input should be a subclass of {class}")]
+    IsSubclassOf {
+        class: String,
+    },
     #[strum(message = "Input should be callable")]
     CallableType,
     // ---------------------
@@ -432,6 +436,7 @@ impl ErrorKind {
             Self::DatetimeObjectInvalid { .. } => extract_context!(DatetimeObjectInvalid, ctx, error: String),
             Self::TimeDeltaParsing { .. } => extract_context!(Cow::Owned, TimeDeltaParsing, ctx, error: String),
             Self::IsInstanceOf { .. } => extract_context!(IsInstanceOf, ctx, class: String),
+            Self::IsSubclassOf { .. } => extract_context!(IsSubclassOf, ctx, class: String),
             Self::UnionTagInvalid { .. } => extract_context!(
                 UnionTagInvalid,
                 ctx,
@@ -520,6 +525,7 @@ impl ErrorKind {
             Self::DatetimeObjectInvalid { error } => render!(self, error),
             Self::TimeDeltaParsing { error } => render!(self, error),
             Self::IsInstanceOf { class } => render!(self, class),
+            Self::IsSubclassOf { class } => render!(self, class),
             Self::UnionTagInvalid {
                 discriminator,
                 tag,
@@ -568,6 +574,7 @@ impl ErrorKind {
             Self::DatetimeObjectInvalid { error } => py_dict!(py, error),
             Self::TimeDeltaParsing { error } => py_dict!(py, error),
             Self::IsInstanceOf { class } => py_dict!(py, class),
+            Self::IsSubclassOf { class } => py_dict!(py, class),
             Self::UnionTagInvalid {
                 discriminator,
                 tag,

src/input/input_abstract.rs

@@ -38,10 +38,6 @@ pub trait Input<'a>: fmt::Debug + ToPyObject {
 
     fn is_none(&self) -> bool;
 
-    fn is_type(&self, _class: &PyType) -> ValResult<bool> {
-        Ok(false)
-    }
-
     #[cfg_attr(has_no_coverage, no_coverage)]
     fn get_attr(&self, _name: &PyString) -> Option<&PyAny> {
         None
@@ -50,6 +46,14 @@ pub trait Input<'a>: fmt::Debug + ToPyObject {
     // input_ prefix to differentiate from the function on PyAny
     fn input_is_instance(&self, class: &PyAny, json_mask: u8) -> PyResult<bool>;
 
+    fn is_exact_instance(&self, _class: &PyType) -> PyResult<bool> {
+        Ok(false)
+    }
+
+    fn input_is_subclass(&self, _class: &PyType) -> PyResult<bool> {
+        Ok(false)
+    }
+
     fn callable(&self) -> bool {
         false
     }

src/input/input_python.rs

@@ -85,10 +85,6 @@ impl<'a> Input<'a> for PyAny {
         self.is_none()
     }
 
-    fn is_type(&self, class: &PyType) -> ValResult<bool> {
-        Ok(self.get_type().eq(class)?)
-    }
-
     fn get_attr(&self, name: &PyString) -> Option<&PyAny> {
         self.getattr(name).ok()
     }
@@ -101,6 +97,18 @@ impl<'a> Input<'a> for PyAny {
         Ok(result == 1)
     }
 
+    fn is_exact_instance(&self, class: &PyType) -> PyResult<bool> {
+        self.get_type().eq(class)
+    }
+
+    fn input_is_subclass(&self, class: &PyType) -> PyResult<bool> {
+        if let Ok(py_type) = self.cast_as::<PyType>() {
+            py_type.is_subclass(class)
+        } else {
+            Ok(false)
+        }
+    }
+
     fn callable(&self) -> bool {
         self.is_callable()
     }

src/validators/is_subclass.rs

@@ -0,0 +1,67 @@
+use pyo3::intern;
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyType};
+
+use crate::build_tools::SchemaDict;
+use crate::errors::{ErrorKind, ValError, ValResult};
+use crate::input::Input;
+use crate::recursion_guard::RecursionGuard;
+
+use super::{BuildContext, BuildValidator, CombinedValidator, Extra, Validator};
+
+#[derive(Debug, Clone)]
+pub struct IsSubclassValidator {
+    class: Py<PyType>,
+    class_repr: String,
+    name: String,
+}
+
+impl BuildValidator for IsSubclassValidator {
+    const EXPECTED_TYPE: &'static str = "is-subclass";
+
+    fn build(
+        schema: &PyDict,
+        _config: Option<&PyDict>,
+        _build_context: &mut BuildContext,
+    ) -> PyResult<CombinedValidator> {
+        let py = schema.py();
+        let class: &PyType = schema.get_as_req(intern!(py, "cls"))?;
+
+        let class_repr = match schema.get_as(intern!(py, "cls_repr"))? {
+            Some(s) => s,
+            None => class.name()?.to_string(),
+        };
+        let name = format!("{}[{}]", Self::EXPECTED_TYPE, class_repr);
+        Ok(Self {
+            class: class.into(),
+            class_repr,
+            name,
+        }
+        .into())
+    }
+}
+
+impl Validator for IsSubclassValidator {
+    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> {
+        match input.input_is_subclass(self.class.as_ref(py))? {
+            true => Ok(input.to_object(py)),
+            false => Err(ValError::new(
+                ErrorKind::IsSubclassOf {
+                    class: self.class_repr.clone(),
+                },
+                input,
+            )),
+        }
+    }
+
+    fn get_name(&self) -> &str {
+        &self.name
+    }
+}

src/validators/mod.rs

@@ -32,6 +32,7 @@ mod function;
 mod generator;
 mod int;
 mod is_instance;
+mod is_subclass;
 mod json;
 mod list;
 mod literal;
@@ -370,6 +371,7 @@ pub fn build_validator<'a>(
         timedelta::TimeDeltaValidator,
         // introspection types
         is_instance::IsInstanceValidator,
+        is_subclass::IsSubclassValidator,
         callable::CallableValidator,
         // arguments
         arguments::ArgumentsValidator,
@@ -488,6 +490,7 @@ pub enum CombinedValidator {
     Timedelta(timedelta::TimeDeltaValidator),
     // introspection types
     IsInstance(is_instance::IsInstanceValidator),
+    IsSubclass(is_subclass::IsSubclassValidator),
     Callable(callable::CallableValidator),
     // arguments
     Arguments(arguments::ArgumentsValidator),

src/validators/new_class.rs

@@ -74,7 +74,7 @@ impl Validator for NewClassValidator {
         recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
         let class = self.class.as_ref(py);
-        if input.is_type(class)? {
+        if input.is_exact_instance(class)? {
             if self.revalidate {
                 let fields_set = input.get_attr(intern!(py, "__fields_set__"));
                 let output = self.validator.validate(py, input, extra, slots, recursion_guard)?;

tests/test_errors.py

@@ -239,6 +239,7 @@ def f(input_value, **kwargs):
         ('time_delta_parsing', 'Input should be a valid timedelta, foobar', {'error': 'foobar'}),
         ('frozen_set_type', 'Input should be a valid frozenset', None),
         ('is_instance_of', 'Input should be an instance of Foo', {'class': 'Foo'}),
+        ('is_subclass_of', 'Input should be a subclass of Foo', {'class': 'Foo'}),
         ('callable_type', 'Input should be callable', None),
         (
             'union_tag_invalid',

tests/validators/test_is_subclass.py

@@ -0,0 +1,79 @@
+import pytest
+
+from pydantic_core import SchemaError, SchemaValidator, ValidationError, core_schema
+
+
+class Foo:
+    pass
+
+
+class Foobar(Foo):
+    pass
+
+
+class Bar:
+    pass
+
+
+def test_is_subclass_basic():
+    v = SchemaValidator(core_schema.is_subclass_schema(Foo))
+    assert v.validate_python(Foo) == Foo
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_python(Bar)
+    # insert_assert(exc_info.value.errors())
+    assert exc_info.value.errors() == [
+        {
+            'kind': 'is_subclass_of',
+            'loc': [],
+            'message': 'Input should be a subclass of Foo',
+            'input_value': Bar,
+            'context': {'class': 'Foo'},
+        }
+    ]
+
+
+@pytest.mark.parametrize(
+    'input_value,valid',
+    [
+        (Foo, True),
+        (Foobar, True),
+        (Bar, False),
+        (type, False),
+        (1, False),
+        ('foo', False),
+        (Foo(), False),
+        (Foobar(), False),
+        (Bar(), False),
+    ],
+)
+def test_is_subclass(input_value, valid):
+    v = SchemaValidator(core_schema.is_subclass_schema(Foo))
+    assert v.isinstance_python(input_value) == valid
+
+
+def test_not_parent():
+    v = SchemaValidator(core_schema.is_subclass_schema(Foobar))
+    assert v.isinstance_python(Foobar)
+    assert not v.isinstance_python(Foo)
+
+
+def test_invalid_type():
+    with pytest.raises(SchemaError, match="TypeError: 'Foo' object cannot be converted to 'PyType"):
+        SchemaValidator(core_schema.is_subclass_schema(Foo()))
+
+
+def test_custom_repr():
+    v = SchemaValidator(core_schema.is_subclass_schema(Foo, cls_repr='Spam'))
+    assert v.validate_python(Foo) == Foo
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_python(Bar)
+    # insert_assert(exc_info.value.errors())
+    assert exc_info.value.errors() == [
+        {
+            'kind': 'is_subclass_of',
+            'loc': [],
+            'message': 'Input should be a subclass of Spam',
+            'input_value': Bar,
+            'context': {'class': 'Spam'},
+        }
+    ]