diff --git a/crates/pyrefly_python/src/module_name.rs b/crates/pyrefly_python/src/module_name.rs index 85d1cabb73..6459b42bd1 100644 --- a/crates/pyrefly_python/src/module_name.rs +++ b/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") } diff --git a/crates/pyrefly_types/src/types.rs b/crates/pyrefly_types/src/types.rs index dfebe448f1..d8ab485866 100644 --- a/crates/pyrefly_types/src/types.rs +++ b/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, + /// The inferred type of the first positional argument when the metadata is a call. + pub first_arg: Option>, +} + #[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, Box<[Type]>), + Annotated(Box, Box<[AnnotatedMetadata]>), Unpack(Box), TypeVar(TypeVar), ParamSpec(ParamSpec), diff --git a/pyrefly/lib/alt/class/class_field.rs b/pyrefly/lib/alt/class/class_field.rs index eedd8656f2..36937efbd9 100644 --- a/pyrefly/lib/alt/class/class_field.rs +++ b/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(|| { diff --git a/pyrefly/lib/alt/class/pydantic.rs b/pyrefly/lib/alt/class/pydantic.rs index 01e8de7646..e55147a983 100644 --- a/pyrefly/lib/alt/class/pydantic.rs +++ b/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=)]` + /// Extract Pydantic field and validator metadata from an `Annotated` annotation. pub fn extract_pydantic_field_from_annotation( &self, annot: Idx, field_name: &Name, + annotated_field_ty: Option<&Type>, metadata: &ClassMetadata, ) -> Option { 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 { + 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, diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index dcb9c0ebb1..8688bfb653 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/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 { diff --git a/pyrefly/lib/alt/specials.rs b/pyrefly/lib/alt/specials.rs index 6917d69c4c..a8de9b6956 100644 --- a/pyrefly/lib/alt/specials.rs +++ b/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 = arguments[1..] + let metadata_types: Vec = 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 = 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( diff --git a/pyrefly/lib/alt/types/class_bases.rs b/pyrefly/lib/alt/types/class_bases.rs index c7b5c9b231..edd6dc1a5f 100644 --- a/pyrefly/lib/alt/types/class_bases.rs +++ b/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 diff --git a/pyrefly/lib/query/type_table.rs b/pyrefly/lib/query/type_table.rs index 32fcde3d06..d63259dfee 100644 --- a/pyrefly/lib/query/type_table.rs +++ b/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()) } diff --git a/pyrefly/lib/test/pydantic/field.rs b/pyrefly/lib/test/pydantic/field.rs index fc955e2752..4d3166518c 100644 --- a/pyrefly/lib/test/pydantic/field.rs +++ b/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#"