fix use all usages within class (or within __init__) to infer type, not just initial assignment #1159

pyrefly/lib/alt/class/class_field.rs

@@ -61,6 +61,7 @@ use crate::binding::binding::ClassFieldDefinition;
 use crate::binding::binding::ExprOrBinding;
 use crate::binding::binding::KeyClassField;
 use crate::binding::binding::KeyClassSynthesizedFields;
+use crate::binding::binding::MethodDefinedAttribute;
 use crate::binding::binding::MethodSelfKind;
 use crate::binding::binding::MethodThatSetsAttr;
 use crate::config::error_kind::ErrorKind;
@@ -1370,6 +1371,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         name: &Name,
         range: TextRange,
         field_definition: &ClassFieldDefinition,
+        method_assignments: &[MethodDefinedAttribute],
         functional_class_def: bool,
         errors: &ErrorCollector,
     ) -> ClassField {
@@ -1393,7 +1395,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         let (
             initialization,
             is_function_without_return_annotation,
-            value_ty,
+            mut value_ty,
             annotation,
             is_inherited,
             direct_annotation,
@@ -1676,6 +1678,35 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             }
         };
 
+        let mut merged_method_assignments = false;
+        if !method_assignments.is_empty()
+            && annotation
+                .as_ref()
+                .and_then(|ann| ann.ty.as_ref())
+                .is_none()
+        {
+            let ignore_errors = self.error_swallower();
+            let mut types = Vec::with_capacity(method_assignments.len() + 1);
+            types.push(value_ty);
+            for assignment in method_assignments {
+                let direct_annotation = assignment
+                    .annotation
+                    .map(|annot| self.get_idx(annot).annotation.clone());
+                let (ty, _, _) = self.analyze_class_field_value(
+                    &assignment.value,
+                    class,
+                    name,
+                    direct_annotation.as_ref(),
+                    true,
+                    assignment.range,
+                    &ignore_errors,
+                );
+                types.push(ty);
+            }
+            value_ty = self.unions(types);
+            merged_method_assignments = true;
+        }
+
         if let Some(annotation) = direct_annotation.as_ref() {
             self.validate_direct_annotation(
                 annotation,
@@ -1692,7 +1723,10 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
 
         // Determine the final type, promoting literals when appropriate.
         let mut has_implicit_literal = value_ty.is_implicit_literal();
-        if !has_implicit_literal && matches!(initialization, ClassFieldInitialization::Method) {
+        if !has_implicit_literal
+            && (matches!(initialization, ClassFieldInitialization::Method)
+                || merged_method_assignments)
+        {
             value_ty.universe(&mut |current_type_node| {
                 has_implicit_literal |= current_type_node.is_implicit_literal();
             });

pyrefly/lib/alt/solve.rs

@@ -2218,6 +2218,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 &field.name,
                 field.range,
                 &field.definition,
+                field.method_assignments.as_ref(),
                 functional_class_def,
                 errors,
             ),

pyrefly/lib/binding/binding.rs

@@ -9,6 +9,7 @@ use std::fmt;
 use std::fmt::Debug;
 use std::fmt::Display;
 use std::hash::Hash;
+use std::sync::Arc;
 
 use dupe::Dupe;
 use pyrefly_derive::TypeEq;
@@ -116,7 +117,7 @@ assert_words!(BindingClassBaseType, 3);
 assert_words!(BindingClassMetadata, 9);
 assert_bytes!(BindingClassMro, 4);
 assert_bytes!(BindingAbstractClassCheck, 4);
-assert_words!(BindingClassField, 21);
+assert_words!(BindingClassField, 23);
 assert_bytes!(BindingClassSynthesizedFields, 4);
 assert_bytes!(BindingLegacyTypeParam, 16);
 assert_words!(BindingYield, 4);
@@ -2581,16 +2582,18 @@ pub struct BindingClassField {
     pub name: Name,
     pub range: TextRange,
     pub definition: ClassFieldDefinition,
+    pub method_assignments: Arc<[MethodDefinedAttribute]>,
 }
 
 impl DisplayWith<Bindings> for BindingClassField {
     fn fmt(&self, f: &mut fmt::Formatter<'_>, ctx: &Bindings) -> fmt::Result {
         write!(
             f,
-            "BindingClassField({}, {}, {})",
+            "BindingClassField({}, {}, {}, method_assignments = {})",
             ctx.display(self.class_idx),
             self.name,
             self.definition.display_with(ctx),
+            self.method_assignments.len(),
         )
     }
 }
@@ -2613,6 +2616,14 @@ pub enum MethodSelfKind {
     Class,
 }
 
+#[derive(Clone, Debug)]
+pub struct MethodDefinedAttribute {
+    pub method: MethodThatSetsAttr,
+    pub value: ExprOrBinding,
+    pub annotation: Option<Idx<KeyAnnotation>>,
+    pub range: TextRange,
+}
+
 /// Bindings for fields synthesized by a class, such as a dataclass's `__init__` method. This
 /// has to be its own key/binding type because of the dependencies between the various pieces of
 /// information about a class: ClassDef -> ClassMetadata -> ClassField -> ClassSynthesizedFields.

pyrefly/lib/binding/class.rs

@@ -6,6 +6,7 @@
  */
 
 use std::mem;
+use std::sync::Arc;
 use std::sync::LazyLock;
 
 use dupe::Dupe as _;
@@ -67,6 +68,7 @@ use crate::binding::binding::KeyExpect;
 use crate::binding::binding::KeyTParams;
 use crate::binding::binding::KeyVariance;
 use crate::binding::binding::KeyVarianceCheck;
+use crate::binding::binding::MethodDefinedAttribute;
 use crate::binding::bindings::BindingsBuilder;
 use crate::binding::bindings::CurrentIdx;
 use crate::binding::bindings::LegacyTParamCollector;
@@ -268,7 +270,8 @@ impl<'a> BindingsBuilder<'a> {
         let mut django_foreign_key_fields: Vec<Name> = Vec::new();
         let mut django_fields_with_choices: Vec<Name> = Vec::new();
         let mut fields = SmallMap::with_capacity(field_definitions.len());
-        for (name, (definition, range)) in field_definitions.into_iter_hashed() {
+        for (name, (definition, range, method_assignments)) in field_definitions.into_iter_hashed()
+        {
             if let ClassFieldDefinition::AssignedInBody {
                 value: ExprOrBinding::Expr(e),
                 ..
@@ -316,6 +319,7 @@ impl<'a> BindingsBuilder<'a> {
                 name: name.into_key(),
                 range,
                 definition,
+                method_assignments: Arc::from(method_assignments.into_boxed_slice()),
             };
             self.insert_binding(key_field, binding);
         }
@@ -676,6 +680,7 @@ impl<'a> BindingsBuilder<'a> {
                     name: member_name,
                     range,
                     definition,
+                    method_assignments: Arc::from([] as [MethodDefinedAttribute; 0]),
                 },
             );
         }

pyrefly/lib/binding/scope.rs

@@ -57,6 +57,7 @@ use crate::binding::binding::KeyVariance;
 use crate::binding::binding::KeyVarianceCheck;
 use crate::binding::binding::KeyYield;
 use crate::binding::binding::KeyYieldFrom;
+use crate::binding::binding::MethodDefinedAttribute;
 use crate::binding::binding::MethodSelfKind;
 use crate::binding::binding::MethodThatSetsAttr;
 use crate::binding::binding::NarrowUseLocation;
@@ -2121,7 +2122,7 @@ impl Scopes {
     /// - Panics if the current scope is not a class body.
     pub fn finish_class_and_get_field_definitions(
         &mut self,
-    ) -> SmallMap<Name, (ClassFieldDefinition, TextRange)> {
+    ) -> SmallMap<Name, (ClassFieldDefinition, TextRange, Vec<MethodDefinedAttribute>)> {
         let mut field_definitions = SmallMap::new();
         let class_body = self.pop();
         let class_scope = {
@@ -2182,26 +2183,54 @@ impl Scopes {
                         definition: value.idx,
                     },
                 };
-                field_definitions.insert_hashed(name.owned(), (definition, static_info.range));
+                field_definitions
+                    .insert_hashed(name.owned(), (definition, static_info.range, Vec::new()));
             }
         });
+        let mut method_defined_fields: SmallMap<Name, Vec<MethodDefinedAttribute>> =
+            SmallMap::new();
         class_scope.method_defined_attributes().for_each(
             |(name, method, InstanceAttribute(value, annotation, range, _))| {
-                if !field_definitions.contains_key_hashed(name.as_ref()) {
-                    field_definitions.insert_hashed(
-                        name,
-                        (
-                            ClassFieldDefinition::DefinedInMethod {
-                                value,
-                                annotation,
-                                method,
-                            },
-                            range,
-                        ),
-                    );
-                }
+                method_defined_fields.entry_hashed(name).or_default().push(
+                    MethodDefinedAttribute {
+                        method,
+                        value,
+                        annotation,
+                        range,
+                    },
+                );
             },
         );
+        method_defined_fields
+            .into_iter_hashed()
+            .for_each(
+                |(name, assignments)| match field_definitions.entry_hashed(name) {
+                    Entry::Occupied(mut occupied) => {
+                        occupied.get_mut().2.extend(assignments);
+                    }
+                    Entry::Vacant(vacant) => {
+                        let mut iter = assignments.into_iter();
+                        if let Some(MethodDefinedAttribute {
+                            method,
+                            value,
+                            annotation,
+                            range,
+                        }) = iter.next()
+                        {
+                            let rest: Vec<MethodDefinedAttribute> = iter.collect();
+                            vacant.insert((
+                                ClassFieldDefinition::DefinedInMethod {
+                                    value,
+                                    annotation,
+                                    method,
+                                },
+                                range,
+                                rest,
+                            ));
+                        }
+                    }
+                },
+            );
         field_definitions
     }
 

pyrefly/lib/test/attributes.rs

@@ -81,6 +81,23 @@ def f(a: A):
     "#,
 );
 
+testcase!(
+    test_infers_attribute_union_from_class_and_method_assignments,
+    r#"
+from typing import assert_type
+
+class A:
+    value = 1
+
+    def promote(self, flag: bool) -> None:
+        if flag:
+            self.value = "a"
+
+def takes(a: A) -> None:
+    assert_type(a.value, int | str)
+    "#,
+);
+
 testcase!(
     test_unannotated_attribute_bad_assignment,
     r#"
@@ -2220,6 +2237,19 @@ def f(a: A):
     "#,
 );
 
+testcase!(
+    test_class_attr_infers_from_method_assignments,
+    r#"
+from typing import assert_type
+class A:
+    a = 1
+    def set_a(self, value: str):
+        self.a = value
+def f(a: A):
+    assert_type(a.a, int | str)
+    "#,
+);
+
 testcase!(
     test_do_not_promote_explicit_literal_param,
     r#"