diff --git a/pyrefly/lib/alt/attr.rs b/pyrefly/lib/alt/attr.rs index efdeebc0de..81e4024d6c 100644 --- a/pyrefly/lib/alt/attr.rs +++ b/pyrefly/lib/alt/attr.rs @@ -299,9 +299,8 @@ pub enum NoAccessReason { /// How to check the value assigned to an attribute on one member of the base. pub enum AttrSetter { - /// Check the value against a declared type. If `narrow` is true, the checked type of the - /// value may be used to narrow the attribute. - Declared { ty: Type, narrow: bool }, + /// Check the value against a declared type and apply the specified narrowing rule. + Declared { ty: Type, narrow: AttrNarrowing }, /// Call `__setattr__` with the value. SetAttr(Type), /// Call a property setter with the value. @@ -310,6 +309,14 @@ pub enum AttrSetter { DescriptorSetter(Type, DescriptorBase), } +/// When an assignment can narrow the attribute's read type. +pub enum AttrNarrowing { + Always, + Never, + /// Narrow only when the assigned value fits the read type. + IfAssignable(Type), +} + #[derive(Debug)] enum InternalError { /// An internal error caused by `as_attribute_base` being partial. @@ -1127,7 +1134,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { got.clone() } }; - narrow.then_some(got_ty) + match narrow { + AttrNarrowing::Always => Some(got_ty), + AttrNarrowing::Never => None, + AttrNarrowing::IfAssignable(read_type) => { + self.is_subset_eq(&got_ty, &read_type).then_some(got_ty) + } + } } AttrSetter::SetAttr(setattr_ty) => { self.call_setattr( @@ -1226,7 +1239,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { Attribute::Simple(attr_ty) => { setters.push(Some(AttrSetter::Declared { ty: attr_ty, - narrow: true, + narrow: AttrNarrowing::Always, })); } Attribute::ClassAttribute(class_attr) => { @@ -1488,8 +1501,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { new_bound_method.func = BoundMethodType::Overload(filtered_overload); let new_method = self.heap.mk_bound_method(new_bound_method); Some(Attribute::ClassAttribute(match class_attr { - ClassAttribute::ReadWrite(_, is_method, is_instance_var) => { - ClassAttribute::ReadWrite(new_method, *is_method, *is_instance_var) + ClassAttribute::ReadWrite(_, is_method, is_instance_var, set_type) => { + ClassAttribute::ReadWrite( + new_method, + *is_method, + *is_instance_var, + set_type.clone(), + ) } ClassAttribute::ReadOnly(_, reason) => { ClassAttribute::ReadOnly(new_method, reason.clone()) @@ -1538,7 +1556,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { { let want_is_dunder_method = attr_name.starts_with("__") && attr_name.ends_with("__") - && matches!(want, ClassAttribute::ReadWrite(_, IsMethod::Yes, _)); + && matches!(want, ClassAttribute::ReadWrite(_, IsMethod::Yes, ..)); for (got_attr, _) in got_attrs.iter() { // A `__getattr__` fallback is not evidence that the missing member exists // when Python would not consult `__getattr__` for it: `__getattr__` governs @@ -1741,6 +1759,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let found_on = found_on.clone(); let mut types = Vec::with_capacity(candidates.len()); + let mut write_types = None; let mut read_only_reason = None; let mut is_class_attribute = false; let mut is_method = IsMethod::No; @@ -1749,29 +1768,36 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let [(attribute, _)] = base_results.as_slice() else { return None; }; - let ty = match attribute { - Attribute::Simple(ty) => ty, + let (ty, set_type) = match attribute { + Attribute::Simple(ty) => (ty, None), Attribute::ClassAttribute(ClassAttribute::ReadWrite( ty, base_is_method, base_is_instance_var, + set_type, )) => { is_class_attribute = true; if *base_is_method == IsMethod::Yes { is_method = IsMethod::Yes; } is_instance_var |= *base_is_instance_var; - ty + (ty, set_type.as_ref()) } Attribute::ClassAttribute(ClassAttribute::ReadOnly(ty, reason)) => { is_class_attribute = true; // A write has to be valid for every base, so the result is read-only if any of // them are. read_only_reason.get_or_insert_with(|| reason.clone()); - ty + (ty, None) } _ => return None, }; + if set_type.is_some() && write_types.is_none() { + write_types = Some(types.clone()); + } + if let Some(write_types) = &mut write_types { + write_types.push(set_type.unwrap_or(ty).clone()); + } types.push(ty.clone()); } let combined = match combine { @@ -1787,6 +1813,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { combined, is_method, is_instance_var, + write_types.map(|types| self.intersects(&types)), )), } } else { diff --git a/pyrefly/lib/alt/class/class_field.rs b/pyrefly/lib/alt/class/class_field.rs index b3f429fbb1..96eaf656a8 100644 --- a/pyrefly/lib/alt/class/class_field.rs +++ b/pyrefly/lib/alt/class/class_field.rs @@ -52,6 +52,7 @@ use starlark_map::small_set::SmallSet; use crate::alt::answers::LookupAnswer; use crate::alt::answers_solver::AnswersSolver; +use crate::alt::attr::AttrNarrowing; use crate::alt::attr::AttrSetter; use crate::alt::attr::AttrSubsetError; use crate::alt::attr::ClassBase; @@ -127,10 +128,10 @@ pub enum IsMethod { /// class access, and possibly through a special case lookup such as a type var with a bound). #[derive(Debug, Clone)] pub enum ClassAttribute { - /// A read-write attribute with a closed form type for both get and set actions. + /// A read-write attribute with a read type and an optional distinct assignment type. /// The `IsMethod` flag enables the `method-assign` diagnostic. /// The boolean marks annotated instance-variable declarations for `instance-var-assign`. - ReadWrite(Type, IsMethod, bool), + ReadWrite(Type, IsMethod, bool, Option), /// A read-only attribute with a closed form type for get actions. ReadOnly(Type, ReadOnlyReason), /// A `NoAccess` attribute indicates that the attribute is well-defined, but does @@ -157,13 +158,13 @@ pub enum ClassAttribute { impl ClassAttribute { pub fn read_write(ty: Type) -> Self { - Self::ReadWrite(ty, IsMethod::No, false) + Self::ReadWrite(ty, IsMethod::No, false, None) } /// Like [`read_write`](Self::read_write) but marks the attribute as a /// class-body-declared instance variable (annotated, not `ClassVar`). pub fn read_write_instance_var(ty: Type) -> Self { - Self::ReadWrite(ty, IsMethod::No, true) + Self::ReadWrite(ty, IsMethod::No, true, None) } pub fn read_only(ty: Type, reason: ReadOnlyReason) -> Self { @@ -405,6 +406,8 @@ enum ClassFieldInner { /// should succeed. ClassAttribute { ty: Type, + /// Django fields can accept values with a different type from their read type. + django_set_type: Option, annotation: Option, initialization: ClassFieldInitialization, read_only_reason: Option, @@ -483,6 +486,7 @@ impl ClassField { Self( ClassFieldInner::ClassAttribute { ty, + django_set_type: None, annotation, initialization, read_only_reason, @@ -596,6 +600,7 @@ impl ClassField { ClassField( ClassFieldInner::ClassAttribute { ty, + django_set_type: None, annotation: None, initialization: ClassFieldInitialization::ClassBody(None), read_only_reason: None, @@ -615,6 +620,7 @@ impl ClassField { Self( ClassFieldInner::ClassAttribute { ty: heap.mk_any_implicit(), + django_set_type: None, annotation: None, initialization: ClassFieldInitialization::recursive(), read_only_reason: None, @@ -709,6 +715,7 @@ impl ClassField { } ClassFieldInner::ClassAttribute { ty, + django_set_type, annotation, initialization, read_only_reason, @@ -721,9 +728,14 @@ impl ClassField { } => { let mut ty = ty.clone(); f(&mut ty); + let mut django_set_type = django_set_type.clone(); + if let Some(ty) = &mut django_set_type { + f(ty); + } Self( ClassFieldInner::ClassAttribute { ty, + django_set_type, annotation: annotation.clone(), initialization: initialization.clone(), read_only_reason: read_only_reason.clone(), @@ -2032,8 +2044,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { has_instance_write: _, } => { let direct_annotation = annot.map(|a| self.get_idx(a).annotation.clone()); - // Check if there's an inherited property or descriptor field from a parent class. - // If so, use the parent's field instead of creating a new field. + // Unannotated Django assignments use the inherited setter. Explicit annotations + // define attributes with normal type and qualifier checks. if !Ast::is_mangled_attr(name) { // Use get_field_from_ancestors to only look at parent classes, not the current class if let Some(parent_field) = self.get_field_from_ancestors( @@ -2042,27 +2054,35 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { name, &|cls, name| self.get_field_from_current_class_only(cls, name), ) { - match parent_field.value.as_ref() { + let inherits_setter = match parent_field.value.as_ref() { ClassField( ClassFieldInner::Property { .. } | ClassFieldInner::Descriptor { .. }, _, - ) => { - // The child is using the inherited property/descriptor setter, - // not redefining the field. Return the parent's field with - // IsInherited::No so the override check skips it. - let ClassField(inner, _) = parent_field.value.into_owned(); - return ClassField(inner, IsInherited::No); - } + ) => true, ClassField( ClassFieldInner::ClassAttribute { - is_instance_var, .. + django_set_type: Some(_), + .. }, _, - ) => inherited_instance_var = *is_instance_var, - _ => { - // For non-property fields, continue with normal processing - } + ) => direct_annotation.is_none(), + _ => false, + }; + if inherits_setter { + // The child uses the inherited setter. IsInherited::No skips + // the override check because the field is not redefined. + let ClassField(inner, _) = parent_field.value.into_owned(); + return ClassField(inner, IsInherited::No); + } + if let ClassField( + ClassFieldInner::ClassAttribute { + is_instance_var, .. + }, + _, + ) = parent_field.value.as_ref() + { + inherited_instance_var = *is_instance_var; } } } @@ -2383,7 +2403,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { false }; - let ty = if let Some(special_ty) = self.get_special_class_field_type( + let (ty, django_set_type) = if let Some(special_types) = self.get_special_class_field_type( class, name, direct_annotation.as_ref(), @@ -2397,9 +2417,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // Don't use the descriptor, since we've set a custom type instead. descriptor_methods = None; descriptor_range = None; - special_ty + special_types } else { - ty + (ty, None) }; // Pin any vars in the type: leaking a var in a class field is particularly @@ -2408,6 +2428,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // TODO(stroxler): Ideally we would implement some simple heuristics, similar to // first-use based inference we use with assignments, to get more useful types here. let ty = self.solver().force(ty); + let django_set_type = django_set_type.map(|ty| self.solver().force(ty)); let descriptor = match (descriptor_methods, &ty) { (Some((getter, setter, deleter)), Type::ClassType(cls)) => Some(Descriptor { range, @@ -2611,6 +2632,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ClassField( ClassFieldInner::ClassAttribute { ty, + django_set_type, annotation, initialization, read_only_reason, @@ -2783,7 +2805,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { /// Apply any class-specific logic for postprocessing the type of a class field. Hook into this /// function if you need to modify the type of an existing field. If you need to add a new - /// field, see ClassSynthesizedField instead. + /// field, see ClassSynthesizedField instead. The result contains the read type and an optional + /// Django assignment type. fn get_special_class_field_type( &self, class: &Class, @@ -2795,7 +2818,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { is_descriptor: bool, range: TextRange, errors: &ErrorCollector, - ) -> Option { + ) -> Option<(Type, Option)> { let initial_value_expr = match field_definition { ClassFieldDefinition::AssignedInBody { value, .. } => { if let ExprOrBinding::Expr(expr) = value.as_ref() { @@ -2819,7 +2842,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ) .or_else(|| self.get_property_class_field_type(class, name, field_definition)) .or_else(|| self.get_pydantic_root_model_class_field_type(class, name)) - .or_else(|| self.get_django_field_type(ty, class, Some(name), initial_value_expr)) + .map(|ty| (ty, None)) + .or_else(|| { + self.get_django_field_types(ty, class, Some(name), initial_value_expr) + .map(|(get_type, set_type)| (get_type, Some(set_type))) + }) } /// Recognize `x = property(fget, fset, fdel)` and return the corresponding @@ -3848,7 +3875,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { if is_classvar { ClassAttribute::read_only(ty, ReadOnlyReason::ClassVar) } else { - ClassAttribute::ReadWrite(ty, IsMethod::Yes, false) + ClassAttribute::ReadWrite(ty, IsMethod::Yes, false, None) } } ClassFieldInner::ProxyMethod { target, .. } => { @@ -3868,6 +3895,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } ClassFieldInner::ClassAttribute { mut ty, + django_set_type, is_classvar, read_only_reason, descriptor_range, @@ -3904,7 +3932,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { if let Some(reason) = read_only_reason { ClassAttribute::read_only(ty, reason) } else { - ClassAttribute::read_write(ty) + ClassAttribute::ReadWrite(ty, IsMethod::No, false, django_set_type) } } } @@ -3984,6 +4012,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { make_bound_classmethod(self.heap, cls, ty).into_inner(), IsMethod::Yes, false, + None, ) } ClassFieldInner::ProxyMethod { .. } => ClassAttribute::no_access( @@ -5777,7 +5806,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { { Some(AttrSetter::Declared { ty: attr_ty, - narrow: true, + narrow: AttrNarrowing::Always, }) } else { errors @@ -5791,7 +5820,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { None } } - ClassAttribute::ReadWrite(attr_ty, is_method, is_instance_var) => { + ClassAttribute::ReadWrite(attr_ty, is_method, is_instance_var, set_type) => { if is_method == IsMethod::Yes { self.error_with_context( errors, @@ -5837,9 +5866,15 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { if let Some(class_object) = class_object { self.check_dataclass_field_final(class_object, attr_name, errors, range); } + if let Some(set_type) = set_type { + return Some(AttrSetter::Declared { + ty: set_type, + narrow: AttrNarrowing::IfAssignable(attr_ty), + }); + } Some(AttrSetter::Declared { ty: attr_ty, - narrow: true, + narrow: AttrNarrowing::Always, }) } ClassAttribute::Property(getter, None, cls) => { @@ -5849,7 +5884,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { if is_cached_property { Some(AttrSetter::Declared { ty: self.call_property_getter(getter, range, errors, context), - narrow: false, + narrow: AttrNarrowing::Never, }) } else { let e = NoAccessReason::SettingReadOnlyProperty(cls); @@ -5895,7 +5930,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // that the value is assignable to the descriptor type. Some(AttrSetter::Declared { ty: self.heap.mk_class_type(x.cls.clone()), - narrow: false, + narrow: AttrNarrowing::Never, }) } }, @@ -6005,8 +6040,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } }; match attr { - ClassAttribute::ReadWrite(ty, is_method, is_instance_var) => Some( - ClassAttribute::ReadWrite(filter_type(ty)?, is_method, is_instance_var), + ClassAttribute::ReadWrite(ty, is_method, is_instance_var, set_type) => Some( + ClassAttribute::ReadWrite(filter_type(ty)?, is_method, is_instance_var, set_type), ), ClassAttribute::ReadOnly(ty, reason) => { Some(ClassAttribute::ReadOnly(filter_type(ty)?, reason)) @@ -6195,7 +6230,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { }) } ( - ClassAttribute::ReadWrite(got, ..), + ClassAttribute::ReadWrite(got, _, _, got_set_type), ClassAttribute::Property(want, want_setter, _), ) => { is_subset( @@ -6221,10 +6256,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { && let Some(rest) = setter_sig.strip_first_param() && let Some(setter_value_type) = rest.get_first_param() { - is_subset(setter_value_type, got).map_err(|subset_error| { + let got_set_type = got_set_type.as_ref().unwrap_or(got); + is_subset(setter_value_type, got_set_type).map_err(|subset_error| { Box::new(AttrSubsetError::Contravariant { want: want_setter.clone(), - got: got.clone(), + got: got_set_type.clone(), got_is_property: false, want_is_property: true, subset_error, diff --git a/pyrefly/lib/alt/class/django.rs b/pyrefly/lib/alt/class/django.rs index 9f719c795c..4bbc46cb29 100644 --- a/pyrefly/lib/alt/class/django.rs +++ b/pyrefly/lib/alt/class/django.rs @@ -47,6 +47,7 @@ use crate::error::collector::ErrorCollector; /// Django stubs use this attribute to specify the Python type that a field should infer to const DJANGO_PRIVATE_GET_TYPE: Name = Name::new_static("_pyi_private_get_type"); +const DJANGO_PRIVATE_SET_TYPE: Name = Name::new_static("_pyi_private_set_type"); pub fn is_django_choices_subclass(bases_with_metadata: &[(Class, &ClassMetadata)]) -> bool { bases_with_metadata.iter().any(|(base, base_meta)| { @@ -247,55 +248,65 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ) } - pub fn get_django_field_type( + /// Return the read and assignment types of a Django field. + pub fn get_django_field_types( &self, ty: &Type, class: &Class, field_name: Option<&Name>, initial_value_expr: Option<&Expr>, - ) -> Option { + ) -> Option<(Type, Type)> { match ty { Type::ClassType(cls) if cls.has_qname(ModuleName::django_utils_functional().as_str(), "_Getter") => { - cls.targs().as_slice().first().cloned() + cls.targs() + .as_slice() + .first() + .map(|ty| (ty.clone(), ty.clone())) } - Type::ClassType(cls) => self.get_django_field_type_from_class( + Type::ClassType(cls) => self.get_django_field_types_from_class( cls.class_object(), class, field_name, initial_value_expr, ), Type::ClassDef(cls) => { - self.get_django_field_type_from_class(cls, class, field_name, initial_value_expr) + self.get_django_field_types_from_class(cls, class, field_name, initial_value_expr) } Type::Union(f) => { - let transformed: Vec<_> = f - .members - .iter() - .map(|variant| { - self.get_django_field_type(variant, class, field_name, initial_value_expr) - .unwrap_or_else(|| variant.clone()) - }) - .collect(); - - if transformed != f.members.to_vec() { - Some(self.unions(transformed)) - } else { - None + let (first_index, (get_type, set_type)) = + f.members.iter().enumerate().find_map(|(index, variant)| { + self.get_django_field_types(variant, class, field_name, initial_value_expr) + .map(|types| (index, types)) + })?; + // Ordinary unions need no collections when none of their members transform. + let mut get_types = Vec::with_capacity(f.members.len()); + let mut set_types = Vec::with_capacity(f.members.len()); + get_types.extend_from_slice(&f.members[..first_index]); + set_types.extend_from_slice(&f.members[..first_index]); + get_types.push(get_type); + set_types.push(set_type); + for variant in &f.members[first_index + 1..] { + let (get_type, set_type) = self + .get_django_field_types(variant, class, field_name, initial_value_expr) + .unwrap_or_else(|| (variant.clone(), variant.clone())); + get_types.push(get_type); + set_types.push(set_type); } + Some((self.unions(get_types), self.unions(set_types))) } _ => None, } } - fn get_django_field_type_from_class( + fn get_django_field_types_from_class( &self, field: &Class, class: &Class, field_name: Option<&Name>, initial_value_expr: Option<&Expr>, - ) -> Option { + ) -> Option<(Type, Type)> { if !(self.get_metadata_for_class(class).is_django_model() && self.inherits_from_django_field(field)) { @@ -309,9 +320,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { && let Some(model_type) = self.resolve_target(to_expr, class) { if self.is_foreign_key_like_field(field) { - Some(model_type) + Some((model_type.clone(), model_type)) } else if self.is_many_to_many_field(field) { - return self.get_manager_type(model_type); + return self.get_manager_type(model_type).map(|ty| (ty.clone(), ty)); } else { None } @@ -320,20 +331,37 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { }; let base_type = base_type.or_else(|| { - self.get_class_member(field, &DJANGO_PRIVATE_GET_TYPE) - .map(|field| field.ty()) + let get_type = self.get_class_member(field, &DJANGO_PRIVATE_GET_TYPE)?.ty(); + let set_type = self.get_class_member(field, &DJANGO_PRIVATE_SET_TYPE)?.ty(); + Some((get_type, set_type)) })?; - let maybe_narrowed_type = - self.narrow_charfield_choices(field, initial_value_expr, base_type); + let (base_get_type, set_type) = base_type; + let get_type = + self.narrow_charfield_choices(field, initial_value_expr, base_get_type.clone()); + let set_type = if get_type != base_get_type { + // Ordinary writes must preserve the finite read type. Expressions retain their + // declared assignment type because their database result is not known here. + let combinable = self.try_get_from_export( + ModuleName::from_str("django.db.models.expressions"), + Name::new_static("Combinable"), + )?; + let combinable = self.class_def_to_instance_type(combinable); + self.union(get_type.clone(), self.intersects(&[set_type, combinable])) + } else { + set_type + }; if let Some(e) = initial_value_expr && let Some(call_expr) = e.as_call_expr() && self.is_django_field_nullable(call_expr) { - Some(self.union(maybe_narrowed_type, self.heap.mk_none())) + Some(( + self.union(get_type, self.heap.mk_none()), + self.union(set_type, self.heap.mk_none()), + )) } else { - Some(maybe_narrowed_type) + Some((get_type, set_type)) } } @@ -786,8 +814,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // No custom pk, use default AutoField type let auto_field_type = self.try_get_from_export(ModuleName::django_models_fields(), AUTO_FIELD)?; - self.get_django_field_type(auto_field_type, model, None, None) - .map(|ty| (ty, false)) + self.get_django_field_types(auto_field_type, model, None, None) + .map(|(get_type, _)| (get_type, false)) } } diff --git a/pyrefly/lib/test/django/fields.rs b/pyrefly/lib/test/django/fields.rs index 4cf83308ab..f3495dfa92 100644 --- a/pyrefly/lib/test/django/fields.rs +++ b/pyrefly/lib/test/django/fields.rs @@ -7,6 +7,287 @@ use crate::django_testcase; +// https://github.com/facebook/pyrefly/issues/5322 +django_testcase!( + test_field_expression_assignment, + r#" +from typing import assert_type +from django.db import models +from django.db.models import F, Value + +class Book(models.Model): + is_draft = models.BooleanField() + pages = models.IntegerField() + title = models.CharField(max_length=100) + optional_pages = models.IntegerField(null=True) + ordinary: bool = False + +def update(book: Book) -> None: + book.is_draft = F("other") + assert_type(book.is_draft, bool) + book.is_draft = Value(True) + assert_type(book.is_draft, bool) + book.pages = F("pages") + 1 + assert_type(book.pages, int) + book.title = Value("New title") + assert_type(book.title, str) + book.optional_pages = F("pages") + assert_type(book.optional_pages, int | None) + book.optional_pages = None + book.is_draft = "invalid" # E: is not assignable to attribute `is_draft` + book.is_draft = None # E: is not assignable to attribute `is_draft` + book.ordinary = F("other") # E: is not assignable to attribute `ordinary` + +class SpecialBook(Book): + pass + +def update_subclass(book: SpecialBook) -> None: + book.is_draft = F("other") + assert_type(book.is_draft, bool) +"#, +); + +django_testcase!( + test_field_assignment_narrowing, + r#" +from typing import assert_type +from django.db import models +from django.db.models import F + +class Author(models.Model): + name = models.CharField(max_length=100) + +class Book(models.Model): + pages = models.IntegerField(null=True) + title = models.CharField(max_length=100, null=True) + author = models.ForeignKey(Author, on_delete=models.CASCADE, null=True) + + def set_pages(self, pages: int) -> int: + self.pages = pages + return self.pages + 1 + +def update(book: Book, pages: int, title: str, author: Author) -> None: + book.pages = pages + assert_type(book.pages, int) + book.title = title + assert_type(book.title, str) + book.author = author + assert_type(book.author, Author) + book.author.name.upper() + book.pages = F("pages") + 1 + assert_type(book.pages, int | None) + book.pages = "123" + assert_type(book.pages, int | None) + book.pages = None + assert_type(book.pages, None) + +def defaults(book: Book) -> int: + if book.pages is None: + book.pages = 1 + if book.title is None: + book.title = "Untitled" + book.title.upper() + return book.pages + 1 + +def update_generic[T: Book](book: T, pages: int) -> int: + book.pages = pages + return book.pages + 1 +"#, +); + +django_testcase!( + test_field_choices_assignment, + r#" +from typing import Literal, assert_type +from django.db import models +from django.db.models import F, Value + +class Card(models.Model): + suit = models.CharField(max_length=1, choices=(("A", "A"), ("B", "B"))) + optional_suit = models.CharField(max_length=1, choices=(("A", "A"), ("B", "B")), null=True) + +def update(card: Card, suit: Literal["A", "B"]) -> None: + card.suit = suit + assert_type(card.suit, Literal["A", "B"]) + card.suit = "INVALID" # E: is not assignable to attribute `suit` + card.suit = 1 # E: is not assignable to attribute `suit` + card.optional_suit = "INVALID" # E: is not assignable to attribute `optional_suit` + card.optional_suit = None + assert_type(card.optional_suit, None) + card.suit = F("suit") + assert_type(card.suit, Literal["A", "B"]) + card.suit = Value("A") + assert_type(card.suit, Literal["A", "B"]) + card.optional_suit = F("suit") + assert_type(card.optional_suit, Literal["A", "B"] | None) +"#, +); + +django_testcase!( + test_field_expression_assignment_receivers, + r#" +from typing import assert_type +from django.db import models +from django.db.models import F, Value + +class Counter(models.Model): + count = models.IntegerField() + optional_count = models.IntegerField(null=True) + + def bump(self) -> None: + self.count = F("count") + 1 + assert_type(self.count, int) + self.optional_count = Value(1) + assert_type(self.optional_count, int | None) + + def bump_explicit(self: Counter) -> None: + self.count = F("count") + 1 + +class SpecialCounter(Counter): + def bump(self) -> None: + self.count = F("count") + 1 + assert_type(self.count, int) + +def update[T: Counter](counter: T) -> None: + counter.count = F("count") + 1 + assert_type(counter.count, int) + counter.optional_count = Value(1) + assert_type(counter.optional_count, int | None) +"#, +); + +django_testcase!( + test_inherited_field_explicit_annotation, + r#" +from typing import Final +from django.db import models + +class Counter(models.Model): + count = models.IntegerField() + +class AnnotatedCounter(Counter): + def reset(self) -> None: + self.count: int = "123" # E: `Literal['123']` is not assignable to attribute `count` with type `int` + +class FinalCounter(Counter): + def __init__(self) -> None: + self.count: Final[int] = 1 # E: `FinalCounter.count` is read-only, but `Counter.count` is read-write + + def reset(self) -> None: + self.count = 2 # E: Cannot set field `count` + +class PlainCounter(Counter): + count: int + + def reset(self) -> None: + self.count = "123" # E: is not assignable to attribute `count` with type `int` +"#, +); + +django_testcase!( + test_field_override_compatibility, + r#" +from django.db import models + +class AbstractName(models.Model): + name = models.CharField(max_length=100) + + class Meta: + abstract = True + +class EmailName(AbstractName): + name = models.EmailField() + +class TextName(AbstractName): + name = models.TextField() + +class AnnotatedName(AbstractName): + name: str + +class PropertyName(AbstractName): + @property + def name(self) -> str: + return "name" + + @name.setter + def name(self, value: str) -> None: + pass + +class BadName(AbstractName): + name: int # E: Class member `BadName.name` overrides parent class `AbstractName` in an inconsistent manner + +class ExplicitSetter(models.Model): + @property + def count(self) -> int: + return 1 + + @count.setter + def count(self, value: str) -> None: + pass + +class FieldSetter(ExplicitSetter): + count = models.IntegerField() + +class BytesSetter(models.Model): + @property + def count(self) -> int: + return 1 + + @count.setter + def count(self, value: bytes) -> None: + pass + +class BadFieldSetter(BytesSetter): + count = models.IntegerField() # E: `BadFieldSetter.count` has type `Combinable | float | int | str`, which is not assignable from `(value: bytes) -> None` +"#, +); + +django_testcase!( + test_field_intersection_assignment, + r#" +from typing import ClassVar, Final +from django.db import models +from django.db.models import F + +class CharRecord(models.Model): + value = models.CharField(max_length=100) + +class TextRecord(models.Model): + value = models.TextField() + +class PlainRecord: + value: str + +class FinalRecord: + value: Final[str] = "fixed" + +class ClassRecord: + value: ClassVar[str] = "shared" + +def update_char(record: CharRecord) -> None: + if isinstance(record, TextRecord): + record.value = "text" + record.value = F("value") + record.value = 1 # E: is not assignable to attribute `value` + if isinstance(record, PlainRecord): + record.value = "text" + record.value = F("value") # E: is not assignable to attribute `value` + if isinstance(record, FinalRecord): + record.value = "text" # E: Cannot set field `value` + if isinstance(record, ClassRecord): + record.value = "text" # E: Cannot set field `value` + +def update_text(record: TextRecord) -> None: + if isinstance(record, CharRecord): + record.value = F("value") + record.value = 1 # E: is not assignable to attribute `value` + +def update_final(record: FinalRecord) -> None: + if isinstance(record, CharRecord): + record.value = "text" # E: Cannot set field `value` +"#, +); + django_testcase!( test_textfield_nullable, r#"