Add a WalkCoreSchema

python/pydantic_core/__init__.py

@@ -21,6 +21,8 @@
     TzInfo,
     Url,
     ValidationError,
+    WalkCoreSchema,
+    WalkCoreSchemaFilterBuilder,
     __version__,
     from_json,
     to_json,
@@ -67,6 +69,8 @@
     'from_json',
     'to_jsonable_python',
     'validate_core_schema',
+    'WalkCoreSchema',
+    'WalkCoreSchemaFilterBuilder',
 ]
 
 

python/pydantic_core/_pydantic_core.pyi

@@ -5,7 +5,7 @@ import sys
 from typing import Any, Callable, Generic, Optional, Type, TypeVar
 
 from pydantic_core import ErrorDetails, ErrorTypeInfo, InitErrorDetails, MultiHostHost
-from pydantic_core.core_schema import CoreConfig, CoreSchema, ErrorType
+from pydantic_core.core_schema import CoreConfig, CoreSchema, ErrorType, SerSchema
 
 if sys.version_info < (3, 8):
     from typing_extensions import final
@@ -46,6 +46,8 @@ __all__ = [
     'list_all_errors',
     'TzInfo',
     'validate_core_schema',
+    'WalkCoreSchema',
+    'WalkCoreSchemaFilterBuilder',
 ]
 __version__: str
 build_profile: str
@@ -864,3 +866,27 @@ def validate_core_schema(schema: CoreSchema, *, strict: bool | None = None) -> C
     We may also remove this function altogether, do not rely on it being present if you are
     using pydantic-core directly.
     """
+
+class _WalkCoreSchemaFilter(Generic[_T]):
+    pass
+
+@final
+class WalkCoreSchemaFilterBuilder(Generic[_T]):
+    def __and__(self, other: WalkCoreSchemaFilterBuilder) -> WalkCoreSchemaFilterBuilder: ...
+    def __or__(self, other: WalkCoreSchemaFilterBuilder) -> WalkCoreSchemaFilterBuilder: ...
+    @staticmethod
+    def has_ref() -> WalkCoreSchemaFilterBuilder: ...
+    @staticmethod
+    def has_type(type: str) -> WalkCoreSchemaFilterBuilder: ...
+    @staticmethod
+    def predicate(predicate: Callable[[_T], bool]) -> WalkCoreSchemaFilterBuilder: ...
+    def build(self, func: Callable[[_T, Callable[[_T], _T]], _T]) -> _WalkCoreSchemaFilter[_T]: ...
+
+@final
+class WalkCoreSchema:
+    def __init__(
+        self,
+        visit_core_schema: _WalkCoreSchemaFilter[CoreSchema] | None = None,
+        visit_ser_schema: _WalkCoreSchemaFilter[SerSchema] | None = None,
+    ) -> None: ...
+    def walk(self, schema: CoreSchema) -> CoreSchema: ...

src/lib.rs

@@ -22,6 +22,7 @@ mod serializers;
 mod tools;
 mod url;
 mod validators;
+mod walk_core_schema;
 
 // required for benchmarks
 pub use self::input::TzInfo;
@@ -111,5 +112,7 @@ fn _pydantic_core(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(to_jsonable_python, m)?)?;
     m.add_function(wrap_pyfunction!(list_all_errors, m)?)?;
     m.add_function(wrap_pyfunction!(validate_core_schema, m)?)?;
+    m.add_class::<walk_core_schema::WalkCoreSchema>()?;
+    m.add_class::<walk_core_schema::WalkCoreSchemaFilterBuilder>()?;
     Ok(())
 }