fix Pydantic union types disable Pyrefly assignment checks #4131

pyrefly/lib/alt/call.rs

@@ -1607,7 +1607,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         if class_metadata.is_pydantic_model()
             && let Some(dataclass) = class_metadata.dataclass_metadata()
         {
-            self.check_pydantic_argument_range_constraints(
+            self.check_pydantic_argument_constraints(
                 cls.class_object(),
                 dataclass,
                 args,

pyrefly/lib/alt/class/pydantic.rs

@@ -105,7 +105,8 @@ impl PydanticRangeConstraints {
 #[derive(Clone)]
 struct PydanticParamConstraint {
     field_name: Name,
-    constraints: PydanticRangeConstraints,
+    range: Option<PydanticRangeConstraints>,
+    lax_type: Option<Type>,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
@@ -975,7 +976,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         None
     }
 
-    pub fn check_pydantic_argument_range_constraints(
+    pub fn check_pydantic_argument_constraints(
         &self,
         cls: &Class,
         dataclass: &DataclassMetadata,
@@ -1196,7 +1197,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                                 &value_ty,
                                 &PydanticParamConstraint {
                                     field_name: name,
-                                    constraints,
+                                    range: Some(constraints),
+                                    lax_type: None,
                                 },
                                 value.range(),
                                 errors,
@@ -1225,18 +1227,31 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         let mut constraints = SmallMap::new();
         let mut position = 0;
         let kw_only_by_class = self.compute_kw_only_fields_by_class(cls);
-        for (field_name, _field, keywords) in
+        for (field_name, field, keywords) in
             self.iter_fields(cls, dataclass, true, &kw_only_by_class)
         {
             if !keywords.init {
                 continue;
             }
-            let Some(constraint) = PydanticRangeConstraints::from_keywords(&keywords) else {
-                continue;
+            let range = PydanticRangeConstraints::from_keywords(&keywords);
+            let lax_type = if !keywords.strict.unwrap_or(dataclass.kws.strict)
+                && !keywords.pydantic_has_custom_schema
+                && keywords.converter_param.is_none()
+                && !dataclass
+                    .pydantic_before_validator_fields
+                    .contains(&field_name)
+            {
+                Some(field.ty())
+            } else {
+                None
             };
+            if range.is_none() && lax_type.is_none() {
+                continue;
+            }
             let info = PydanticParamConstraint {
                 field_name: field_name.clone(),
-                constraints: constraint,
+                range,
+                lax_type,
             };
             if keywords.init_by_name {
                 constraints.insert(PydanticParamKey::Name(field_name), info.clone());
@@ -1269,14 +1284,34 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             }
             return;
         }
+        // The signature checks input types. This check enforces value restrictions on literals.
+        if let Some(field_ty) = &info.lax_type
+            && matches!(value_ty, Type::Literal(_))
+            && self.is_subset_eq(value_ty, &self.expand_type_for_lax_mode(field_ty))
+            && !self.pydantic_lax_accepts_literal(field_ty, value_ty)
+        {
+            self.error(
+                errors,
+                range,
+                ErrorKind::BadArgumentType,
+                format!(
+                    "Argument value `{}` is not valid for Pydantic `bool` field `{}`",
+                    self.for_display(value_ty.clone()),
+                    info.field_name
+                ),
+            );
+        }
+        let Some(constraints) = &info.range else {
+            return;
+        };
         let Some(value_lit) = int_literal_from_type(value_ty) else {
             return;
         };
         let checks = [
-            ("gt", info.constraints.gt.as_ref()),
-            ("ge", info.constraints.ge.as_ref()),
-            ("lt", info.constraints.lt.as_ref()),
-            ("le", info.constraints.le.as_ref()),
+            ("gt", constraints.gt.as_ref()),
+            ("ge", constraints.ge.as_ref()),
+            ("lt", constraints.lt.as_ref()),
+            ("le", constraints.le.as_ref()),
         ];
         for (label, constraint_ty) in checks {
             let Some(constraint_ty) = constraint_ty else {

pyrefly/lib/alt/class/pydantic_lax.rs

@@ -9,12 +9,15 @@ use pyrefly_python::module_name::ModuleName;
 use pyrefly_types::class::ClassType;
 use pyrefly_types::identity::IdentityIgnored;
 use pyrefly_types::keywords::ConverterMap;
+use pyrefly_types::literal::Lit;
 use pyrefly_types::tuple::Tuple;
 use ruff_python_ast::name::Name;
 use starlark_map::ordered_map::OrderedMap;
 
 use crate::alt::answers::LookupAnswer;
 use crate::alt::answers_solver::AnswersSolver;
+use crate::alt::types::class_metadata::DataclassKind;
+use crate::binding::pydantic::GET_PYDANTIC_CORE_SCHEMA;
 use crate::types::class::Class;
 use crate::types::types::Type;
 
@@ -135,6 +138,38 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         }
     }
 
+    /// Check each union alternative with its lax conversion and refine bool's literal inputs.
+    pub fn pydantic_lax_accepts_literal(&self, expected: &Type, literal: &Type) -> bool {
+        match expected {
+            Type::Union(union) => union
+                .members
+                .iter()
+                .any(|member| self.pydantic_lax_accepts_literal(member, literal)),
+            Type::ClassType(cls) if cls == self.stdlib.bool() => match literal {
+                Type::Literal(lit) => match &lit.value {
+                    Lit::Int(value) => matches!(value.as_i64(), Some(0 | 1)),
+                    Lit::Str(value) => matches!(
+                        value.to_ascii_lowercase().as_str(),
+                        "0" | "1"
+                            | "f"
+                            | "false"
+                            | "n"
+                            | "no"
+                            | "off"
+                            | "on"
+                            | "t"
+                            | "true"
+                            | "y"
+                            | "yes"
+                    ),
+                    _ => self.is_subset_eq(literal, &self.expand_type_for_lax_mode(expected)),
+                },
+                _ => self.is_subset_eq(literal, &self.expand_type_for_lax_mode(expected)),
+            },
+            _ => self.is_subset_eq(literal, &self.expand_type_for_lax_mode(expected)),
+        }
+    }
+
     fn get_container_lax_conversion(
         &self,
         class_obj: &Class,
@@ -189,7 +224,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         }
     }
 
-    fn expand_type_for_lax_mode(&self, ty: &Type) -> Type {
+    pub fn expand_type_for_lax_mode(&self, ty: &Type) -> Type {
         match ty {
             Type::None => ty.clone(),
             // Literal types have no lax coercion - they require exact values
@@ -200,10 +235,47 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             Type::Tuple(tuple) => self
                 .heap
                 .mk_class_type(self.stdlib.iterable(self.get_tuple_element_type(tuple))),
-            // Container types: recursively expand all type arguments
-            Type::ClassType(cls) if !cls.targs().as_slice().is_empty() => {
+            Type::ClassType(cls) => {
+                if let Some(converted) = self.get_atomic_lax_conversion(ty) {
+                    return converted;
+                }
                 let class_obj = cls.class_object();
+                // BaseModel's built-in hook preserves model validation. Other hooks may accept any input.
+                if self
+                    .get_non_synthesized_class_member_and_defining_class(
+                        class_obj,
+                        &GET_PYDANTIC_CORE_SCHEMA,
+                    )
+                    .is_some_and(|member| {
+                        !member
+                            .defining_class
+                            .has_toplevel_qname(ModuleName::pydantic().as_str(), "BaseModel")
+                    })
+                {
+                    return self.heap.mk_any_explicit();
+                }
+                let metadata = self.get_metadata_for_class(class_obj);
+                if metadata.is_pydantic_model()
+                    || metadata.dataclass_metadata().is_some_and(|dataclass| {
+                        matches!(
+                            dataclass.kind,
+                            DataclassKind::Dataclass {
+                                is_stdlib_dataclass: true,
+                                ..
+                            }
+                        )
+                    })
+                {
+                    let mapping = self.heap.mk_class_type(self.stdlib.mapping(
+                        self.heap.mk_class_type(self.stdlib.str().clone()),
+                        self.heap.mk_any_explicit(),
+                    ));
+                    return self.union(ty.clone(), mapping);
+                }
                 let targs = cls.targs().as_slice();
+                if targs.is_empty() {
+                    return ty.clone();
+                }
 
                 // Special handling for dict and Mapping: don't expand key type
                 // (Mapping is invariant in key, so expanding it would make
@@ -244,10 +316,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 let expanded_members = self.expand_types(&f.members);
                 self.unions(expanded_members)
             }
-            // Known atomic types with conversion tables, or Any for everything else
-            _ => self
-                .get_atomic_lax_conversion(ty)
-                .unwrap_or_else(|| self.heap.mk_any_explicit()),
+            // Runtime validation determines whether unsupported forms such as callables are valid.
+            _ => self.heap.mk_any_explicit(),
         }
     }
 

pyrefly/lib/test/pydantic/strictness.rs

@@ -304,6 +304,107 @@ reveal_type(Model.__init__)  # E: revealed type: (self: Model, *, y: Decimal | b
     "#,
 );
 
+pydantic_testcase!(
+    test_lax_mode_coercion_union_rejects_invalid_values,
+    r#"
+from typing import Literal
+from pydantic import BaseModel
+
+class Something(BaseModel):
+    pass
+
+class OtherThing(BaseModel):
+    pass
+
+class Model(BaseModel):
+    scalar: str | bool
+    model: Something | OtherThing
+
+Model(scalar=5, model=Something())  # E: Argument value `Literal[5]` is not valid for Pydantic `bool` field `scalar`
+Model(scalar="value", model=5)  # E: Argument `Literal[5]` is not assignable to parameter `model`
+
+class BoolModel(BaseModel):
+    value: bool
+
+BoolModel(value=0)
+BoolModel(value="yes")
+BoolModel(value=2)  # E: Argument value `Literal[2]` is not valid for Pydantic `bool` field `value`
+BoolModel(value="maybe")  # E: Argument value `Literal['maybe']` is not valid for Pydantic `bool` field `value`
+
+class IntOrBoolModel(BaseModel):
+    value: int | bool
+
+IntOrBoolModel(value=5)
+IntOrBoolModel(value="123")
+
+def union_input(value: Literal[0, 5]) -> None:
+    BoolModel(value=value)  # E: Argument value `Literal[5]` is not valid
+    "#,
+);
+
+pydantic_testcase!(
+    test_lax_mode_nested_model_and_dataclass_inputs,
+    r#"
+from dataclasses import dataclass
+from typing import Any, Mapping, assert_type
+from pydantic import BaseModel
+
+class Link(BaseModel):
+    href: str
+
+@dataclass
+class Record:
+    value: int
+
+class Publication(BaseModel):
+    links: list[Link]
+    data: Link | Record
+
+links: list[dict[str, str]] = [{"href": "example"}]
+record: Mapping[str, Any] = {"value": 1}
+publication = Publication(links=links, data=record)
+Publication(links=[Link(href="example")], data=Record(1))
+Publication(links=links, data=Link(href="example"))
+Publication(links=[5], data=5)  # E: is not assignable to parameter `links` # E: is not assignable to parameter `data`
+assert_type(publication.links, list[Link])
+assert_type(publication.data, Link | Record)
+    "#,
+);
+
+pydantic_testcase!(
+    test_lax_mode_custom_schema_inputs,
+    r#"
+from typing import Any, assert_type
+from pydantic import BaseModel, ConfigDict
+
+class CustomSchema:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: Any) -> Any: ...
+
+class InheritedSchema(CustomSchema):
+    pass
+
+class CustomModel(BaseModel):
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: Any) -> Any: ...
+
+class Plain:
+    pass
+
+class Model(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    pa_schema: InheritedSchema | bool
+    custom_model: CustomModel
+    plain: Plain | None
+
+m = Model(pa_schema=InheritedSchema, custom_model=5, plain=Plain())
+Model(pa_schema=5, custom_model="input", plain=None)
+Model(pa_schema=InheritedSchema, custom_model=5, plain=5)  # E: is not assignable to parameter `plain`
+assert_type(m.pa_schema, InheritedSchema | bool)
+assert_type(m.custom_model, CustomModel)
+    "#,
+);
+
 pydantic_testcase!(
     test_lax_mode_list_and_set_invariance,
     r#"

website/docs/pydantic-lax-conversions.mdx

@@ -14,7 +14,7 @@ description: Complete reference of how Pyrefly converts types in Pydantic lax mo
 
 This page provides a complete reference for how Pyrefly converts types when working with Pydantic models in **lax mode** (the default). For background on how lax mode works, see the [main Pydantic documentation](../pydantic).
 
-**Note:** Types without a specific conversion rule (e.g., `Callable`, `Any`, custom classes, and generic classes not listed below) are converted to `Any`.
+**Note:** Nested Pydantic models and standard-library dataclasses accept either an instance or a `Mapping[str, Any]`. Classes with a custom `__get_pydantic_core_schema__` hook convert to `Any`. Other custom class types without a conversion rule are left unchanged. Unsupported forms, such as `Callable` and generic classes not listed below, convert to `Any`.
 
 ---
 
@@ -38,6 +38,8 @@ Named unions are used for atomic types to keep type signatures concise.
 | `UUID` | `LaxUuid` | `UUID \| str` |
 | `None` | (no conversion) | `None` |
 
+For `bool` conversions, Pyrefly checks inferred integer and string literals against Pydantic's accepted values. Non-literal values use the broader input types shown above.
+
 ---
 
 ## Compositional Type Conversions
@@ -61,4 +63,3 @@ Named unions are used for atomic types to keep type signatures concise.
 - **Single-element containers and unbounded tuples:** Named unions are preserved. `list[int]` → `Iterable[LaxInt]`
 - **Concrete tuples:** Element types are expanded and flattened. `tuple[int, str]` → `Iterable[int | bool | float | str | bytes | bytearray | Decimal]`
 - **Dictionaries:** Only values are converted; keys remain unchanged. `dict[str, int]` → `Mapping[str, LaxInt]`
-