Make ModelWrapValidator protocols generic (#7154)

pydantic · Aug 22, 2023 · 1bb2227 · 1bb2227
1 parent 1ba4463
commit 1bb2227
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pydantic/functional_validators.py b/pydantic/functional_validators.py
@@ -351,7 +351,7 @@ def __call__(  # noqa: D102
         ...
 
 
-class ModelWrapValidatorWithoutInfo(Protocol):
+class ModelWrapValidatorWithoutInfo(Protocol[_ModelType]):
     """A @model_validator decorated function signature.
     This is used when `mode='wrap'` and the function does not have info argument.
     """
@@ -368,7 +368,7 @@ def __call__(  # noqa: D102
         ...
 
 
-class ModelWrapValidator(Protocol):
+class ModelWrapValidator(Protocol[_ModelType]):
     """A @model_validator decorated function signature. This is used when `mode='wrap'`."""
 
     def __call__(  # noqa: D102
@@ -423,7 +423,7 @@ def __call__(  # noqa: D102
 ModelAfterValidator = Callable[[_ModelType, _core_schema.ValidationInfo], _ModelType]
 """A `@model_validator` decorated function signature. This is used when `mode='after'`."""
 
-_AnyModelWrapValidator = Union[ModelWrapValidator, ModelWrapValidatorWithoutInfo]
+_AnyModelWrapValidator = Union[ModelWrapValidator[_ModelType], ModelWrapValidatorWithoutInfo[_ModelType]]
 _AnyModeBeforeValidator = Union[ModelBeforeValidator, ModelBeforeValidatorWithoutInfo]
 _AnyModelAfterValidator = Union[ModelAfterValidator[_ModelType], ModelAfterValidatorWithoutInfo[_ModelType]]