fix Support Pydantic validators in Annotated fields #4256

crates/pyrefly_python/src/module_name.rs

@@ -273,6 +273,10 @@ impl ModuleName {
         Self::from_str("pydantic.alias_generators")
     }
 
+    pub fn pydantic_functional_validators() -> Self {
+        Self::from_str("pydantic.functional_validators")
+    }
+
     pub fn django_models_enums() -> Self {
         Self::from_str("django.db.models.enums")
     }

crates/pyrefly_types/src/types.rs

@@ -99,6 +99,17 @@ impl Var {
     }
 }
 
+/// Semantic information retained for an item in `Annotated[T, ...]`.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, TypeEq)]
+pub struct AnnotatedMetadata {
+    /// The inferred type of the metadata value.
+    pub ty: Type,
+    /// The callee's final name when the metadata is a call.
+    pub callee_name: Option<Name>,
+    /// The inferred type of the first positional argument when the metadata is a call.
+    pub first_arg: Option<Box<Type>>,
+}
+
 #[derive(PartialEq, Eq)]
 pub enum TParamsSource {
     Class,
@@ -1012,7 +1023,7 @@ pub enum Type {
     /// This is transparent when resolving annotations, but is not callable and
     /// cannot be assigned to `type[T]`.
     /// The second field carries the metadata items (the `...` in `Annotated[T, ...]`).
-    Annotated(Box<Type>, Box<[Type]>),
+    Annotated(Box<Type>, Box<[AnnotatedMetadata]>),
     Unpack(Box<Type>),
     TypeVar(TypeVar),
     ParamSpec(ParamSpec),

pyrefly/lib/alt/class/class_field.rs

@@ -1722,9 +1722,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                         .is_some_and(|annot| annot.has_qualifier(&Qualifier::ClassVar))
                 {
                     ClassFieldInitialization::Magic
-                } else if let Some(flags) =
-                    self.extract_pydantic_field_from_annotation(*annot, name, metadata)
-                {
+                } else if let Some(flags) = self.extract_pydantic_field_from_annotation(
+                    *annot,
+                    name,
+                    direct_annotation
+                        .as_ref()
+                        .and_then(|annotation| annotation.ty.as_ref()),
+                    metadata,
+                ) {
                     ClassFieldInitialization::ClassBody(Some(Box::new(flags)))
                 } else {
                     ClassFieldInitialization::Uninitialized
@@ -1810,7 +1815,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 let mut direct_annotation = annot.map(|a| self.get_idx(a).annotation.clone());
                 let mut annotation_flags = annot
                     .and_then(|annot| {
-                        self.extract_pydantic_field_from_annotation(annot, name, metadata)
+                        self.extract_pydantic_field_from_annotation(
+                            annot,
+                            name,
+                            direct_annotation
+                                .as_ref()
+                                .and_then(|annotation| annotation.ty.as_ref()),
+                            metadata,
+                        )
                     })
                     .map(|flags| DataclassFieldKeywords {
                         default: flags.default,
@@ -1819,6 +1831,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                         pydantic_unchecked_metadata: flags.pydantic_unchecked_metadata
                             || flags.init_by_alias.is_some(),
                         pydantic_has_custom_schema: flags.pydantic_has_custom_schema,
+                        converter_param: flags.converter_param,
                         ..DataclassFieldKeywords::new()
                     });
                 if metadata.is_protocol()
@@ -1835,7 +1848,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                         ),
                     );
                 }
-                let flags = if let ExprOrBinding::Expr(e) = value.as_ref()
+                let mut flags = if let ExprOrBinding::Expr(e) = value.as_ref()
                     && let Some(dm) = metadata.dataclass_metadata()
                     && let Expr::Call(call) = e
                 {
@@ -1974,6 +1987,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 } else {
                     None
                 };
+                if let Some(flags) = &mut flags
+                    && let Some(annotation_flags) = &mut annotation_flags
+                {
+                    flags.converter_param = annotation_flags.converter_param.take();
+                }
                 let initialization = ClassFieldInitialization::ClassBody(
                     flags
                         .or_else(|| {

pyrefly/lib/alt/class/pydantic.rs

@@ -66,6 +66,7 @@ use crate::binding::pydantic::VALIDATE_BY_ALIAS;
 use crate::binding::pydantic::VALIDATE_BY_NAME;
 use crate::error::collector::ErrorCollector;
 use crate::types::class::Class;
+use crate::types::types::Forallable;
 use crate::types::types::Type;
 
 fn int_literal_from_type(ty: &Type) -> Option<&LitInt> {
@@ -869,13 +870,12 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         }
     }
 
-    /// Extract Pydantic Field metadata from an annotation binding.
-    /// This handles the Pydantic-specific pattern where fields can be declared as:
-    /// `field: Annotated[some_type, Field(some_keyword=<value>)]`
+    /// Extract Pydantic field and validator metadata from an `Annotated` annotation.
     pub fn extract_pydantic_field_from_annotation(
         &self,
         annot: Idx<KeyAnnotation>,
         field_name: &Name,
+        annotated_field_ty: Option<&Type>,
         metadata: &ClassMetadata,
     ) -> Option<DataclassFieldKeywords> {
         let dm = metadata.dataclass_metadata()?;
@@ -963,18 +963,71 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     .get_or_insert_with(DataclassFieldKeywords::new)
                     .strict = Some(strict);
             }
-
             if has_custom_schema {
                 keywords
                     .get_or_insert_with(DataclassFieldKeywords::new)
                     .pydantic_has_custom_schema = true;
             }
 
+            let annotation_ty = self.expr_infer(annotation_expr, &ignore_errors);
+            if let Some(converter_param) =
+                self.pydantic_annotated_validator_input(&annotation_ty, annotated_field_ty)
+            {
+                keywords
+                    .get_or_insert_with(DataclassFieldKeywords::new)
+                    .converter_param = Some(converter_param);
+            }
             return keywords;
         }
         None
     }
 
+    fn pydantic_annotated_validator_input(
+        &self,
+        ty: &Type,
+        annotated_field_ty: Option<&Type>,
+    ) -> Option<Type> {
+        match ty {
+            // Before validators run right-to-left, so the rightmost one receives the raw input.
+            Type::Annotated(_, metadata) => metadata.iter().rev().find_map(|metadata| {
+                let is_validator_name = metadata.callee_name.as_ref().is_some_and(|name| {
+                    matches!(name.as_str(), "BeforeValidator" | "PlainValidator")
+                });
+                let is_validator_type = match &metadata.ty {
+                    Type::ClassType(cls) => {
+                        cls.has_qname(
+                            ModuleName::pydantic_functional_validators().as_str(),
+                            "BeforeValidator",
+                        ) || cls.has_qname(
+                            ModuleName::pydantic_functional_validators().as_str(),
+                            "PlainValidator",
+                        )
+                    }
+                    _ => false,
+                };
+                if !is_validator_name && !is_validator_type {
+                    return None;
+                }
+                metadata
+                    .first_arg
+                    .as_deref()
+                    .map(|func| self.get_converter_param(func, annotated_field_ty))
+            }),
+            Type::TypeAlias(alias) => {
+                let alias = self.get_type_alias(alias);
+                self.pydantic_annotated_validator_input(&alias.as_type(), annotated_field_ty)
+            }
+            Type::Forall(forall) => match &forall.body {
+                Forallable::TypeAlias(alias) => {
+                    let alias = self.get_type_alias(alias);
+                    self.pydantic_annotated_validator_input(&alias.as_type(), annotated_field_ty)
+                }
+                _ => None,
+            },
+            _ => None,
+        }
+    }
+
     pub fn check_pydantic_argument_range_constraints(
         &self,
         cls: &Class,

pyrefly/lib/alt/solve.rs

@@ -189,6 +189,7 @@ use crate::types::type_var::Restriction;
 use crate::types::type_var::TypeVar;
 use crate::types::type_var::Variance;
 use crate::types::type_var_tuple::TypeVarTuple;
+use crate::types::types::AnnotatedMetadata;
 use crate::types::types::AnyStyle;
 use crate::types::types::BoundMethodType;
 use crate::types::types::Forallable;
@@ -1633,7 +1634,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 .bindings()
                 .shape_declarations
                 .shaped_annotation_metadata(range)
-                .map(|shape| vec![Lit::Str(shape.into()).to_implicit_type()].into_boxed_slice()),
+                .map(|shape| {
+                    vec![AnnotatedMetadata {
+                        ty: Lit::Str(shape.into()).to_implicit_type(),
+                        callee_name: None,
+                        first_arg: None,
+                    }]
+                    .into_boxed_slice()
+                }),
         };
         let untyped = self.untype_opt(ty.clone(), range, errors);
         let ty = if let Some(untyped) = untyped {

pyrefly/lib/alt/specials.rs

@@ -32,6 +32,7 @@ use crate::types::lit_int::LitInt;
 use crate::types::literal::Lit;
 use crate::types::special_form::SpecialForm;
 use crate::types::tuple::Tuple;
+use crate::types::types::AnnotatedMetadata;
 use crate::types::types::AnyStyle;
 use crate::types::types::Type;
 
@@ -591,13 +592,31 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             ),
             SpecialForm::Annotated if arguments.len() > 1 => {
                 let mut inner = self.expr_untype(&arguments[0], type_argument_context, errors);
-                let metadata: Vec<Type> = arguments[1..]
+                let metadata_types: Vec<Type> = arguments[1..]
                     .iter()
                     .map(|e| self.expr_infer(e, &self.error_swallower()))
                     .collect();
-                if let Some(frame) = self.polars_frame_annotated_type(&inner, &metadata) {
+                if let Some(frame) = self.polars_frame_annotated_type(&inner, &metadata_types) {
                     inner = frame;
                 }
+                let metadata: Vec<AnnotatedMetadata> = arguments[1..]
+                    .iter()
+                    .zip(metadata_types)
+                    .map(|(e, ty)| AnnotatedMetadata {
+                        ty,
+                        callee_name: e.as_call_expr().and_then(|call| match &*call.func {
+                            Expr::Name(name) => Some(name.id.clone()),
+                            Expr::Attribute(attribute) => Some(attribute.attr.id.clone()),
+                            _ => None,
+                        }),
+                        first_arg: e.as_call_expr().and_then(|call| {
+                            call.arguments.args.first().map(|arg| {
+                                let errors = self.error_swallower();
+                                Box::new(self.expr_infer(arg, &errors))
+                            })
+                        }),
+                    })
+                    .collect();
                 Type::Annotated(Box::new(inner), metadata.into_boxed_slice())
             }
             SpecialForm::Annotated => self.error(

pyrefly/lib/alt/types/class_bases.rs

@@ -286,7 +286,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             if let Type::Annotated(_, metadata) = alias.as_type() {
                 return metadata
                     .iter()
-                    .any(|metadata| self.is_pydantic_strict_metadata(metadata));
+                    .any(|metadata| self.is_pydantic_strict_metadata(&metadata.ty));
             }
         }
         false

pyrefly/lib/query/type_table.rs

@@ -389,7 +389,7 @@ pub(super) fn type_to_indexed_shape(
             args.extend(
                 metadata
                     .iter()
-                    .map(|ty| type_to_indexed_shape(context, ty, table)),
+                    .map(|metadata| type_to_indexed_shape(context, &metadata.ty, table)),
             );
             insert_indexed_named(table, "typing.Annotated", args, None, Vec::new())
         }

pyrefly/lib/test/pydantic/field.rs

@@ -1547,6 +1547,44 @@ P(x=42)
     "#,
 );
 
+pydantic_testcase!(
+    test_annotated_before_validator,
+    r#"
+from typing import Annotated, TypeAlias, reveal_type
+from pydantic import BaseModel, BeforeValidator
+
+def join_ints(values: list[int]) -> str:
+    return ",".join(str(value) for value in values)
+
+def split_ints(value: str) -> list[int]:
+    return [int(part) for part in value.split(",")]
+
+ValidatedStr: TypeAlias = Annotated[str, BeforeValidator(join_ints)]
+ChainedStr: TypeAlias = Annotated[
+    str,
+    BeforeValidator(join_ints),
+    BeforeValidator(split_ints),
+]
+
+class Model(BaseModel):
+    direct: Annotated[str, BeforeValidator(join_ints)]
+    aliased: ValidatedStr
+    chained: ChainedStr
+    defaulted: ValidatedStr = "not validated by default"
+
+model = Model(direct=[1, 2], aliased=[3, 4], chained="5,6", defaulted=[7])
+reveal_type(model.direct)  # E: revealed type: str
+reveal_type(model.aliased)  # E: revealed type: str
+reveal_type(model.chained)  # E: revealed type: str
+reveal_type(model.defaulted)  # E: revealed type: str
+Model(direct=[1, 2], aliased=[3, 4], chained="5,6")
+Model(direct="1,2", aliased=[3, 4], chained="5,6")  # E: Argument `Literal['1,2']` is not assignable to parameter `direct` with type `list[int]` in function `Model.__init__`
+Model(direct=[1, 2], aliased="3,4", chained="5,6")  # E: Argument `Literal['3,4']` is not assignable to parameter `aliased` with type `list[int]` in function `Model.__init__`
+Model(direct=[1, 2], aliased=[3, 4], chained=[5, 6])  # E: Argument `list[int]` is not assignable to parameter `chained` with type `str` in function `Model.__init__`
+Model(direct=[1, 2], aliased=[3, 4], chained="5,6", defaulted="7")  # E: Argument `Literal['7']` is not assignable to parameter `defaulted` with type `list[int]` in function `Model.__init__`
+    "#,
+);
+
 pydantic_testcase!(
     test_field_without_annotation,
     r#"