Add support for descriptor checks with unions

crates/ty_python_semantic/resources/mdtest/descriptor_protocol.md

@@ -194,6 +194,99 @@ reveal_type(C2().attr)  # revealed: Unknown | Literal["non-data", "normal"]
 C2().attr = 1
 ```
 
+### Union of descriptor and non-descriptor types
+
+When an attribute is typed as a union where some elements are data descriptors and some are not,
+assignments should validate against the descriptor's `__set__` method for the descriptor case:
+
+```py
+from typing import Any
+
+class DescriptorWithSet:
+    def __set__(self, instance: object, value: Any) -> None: ...
+
+class NoSet:
+    pass
+
+class MyModel:
+    state: DescriptorWithSet | NoSet = DescriptorWithSet()
+
+m = MyModel()
+# This is valid because `DescriptorWithSet.__set__` accepts `Any`
+m.state = 1
+m.state = "hello"
+```
+
+When the descriptor's `__set__` method has a more restrictive type, only compatible values are
+allowed:
+
+```py
+from typing import Any
+
+class IntDescriptor:
+    def __set__(self, instance: object, value: int) -> None: ...
+
+class NoSet:
+    pass
+
+class MyModel:
+    state: IntDescriptor | NoSet = IntDescriptor()
+
+m = MyModel()
+m.state = 1
+# error: [invalid-assignment] "Invalid assignment to data descriptor attribute `state` on type `MyModel` with custom `__set__` method"
+m.state = "hello"
+```
+
+Multiple descriptors in a union where both have `__set__`:
+
+```py
+from typing import Any
+
+class IntDescriptor:
+    def __set__(self, instance: object, value: int) -> None: ...
+
+class StrDescriptor:
+    def __set__(self, instance: object, value: str) -> None: ...
+
+class MyModel:
+    state: IntDescriptor | StrDescriptor = IntDescriptor()
+
+m = MyModel()
+# For a union of descriptors, we call __set__ on the union of elements that have __set__.
+# The value must be acceptable to all descriptors in the union.
+# error: [invalid-assignment]
+m.state = 1  # int is accepted by IntDescriptor but not StrDescriptor
+# error: [invalid-assignment]
+m.state = "hello"  # str is accepted by StrDescriptor but not IntDescriptor
+```
+
+When a class has both a union descriptor class attribute and an instance annotation, the instance
+annotation takes precedence for the non-descriptor elements:
+
+```py
+from typing import Any
+
+class DescriptorWithSet:
+    def __set__(self, instance: object, value: Any) -> None: ...
+
+class NoSet:
+    pass
+
+class MyModel:
+    state: DescriptorWithSet | NoSet = DescriptorWithSet()
+    def __init__(self):
+        self.state: int
+
+m = MyModel()
+# For descriptor elements, __set__ is called (which accepts Any).
+# For non-descriptor elements, the value goes to instance dict, checked against `int`.
+m.state = 1
+
+# error: [invalid-assignment] "Object of type `Literal["hello"]` is not assignable to attribute `state` of type `int`"
+m.state = "hello"
+```
+
 ### Descriptors only work when used as class variables
 
 Descriptors only work when used as class variables. When put in instances, they have no effect.

crates/ty_python_semantic/src/types/class.rs

@@ -3830,10 +3830,21 @@ impl<'db> ClassLiteral<'db> {
                     if has_binding {
                         // The attribute is declared and bound in the class body.
 
-                        if let Some(implicit_ty) =
-                            Self::implicit_attribute(db, body_scope, name, MethodDecorator::None)
-                                .ignore_possibly_undefined()
+                        let implicit_member =
+                            Self::implicit_attribute(db, body_scope, name, MethodDecorator::None);
+
+                        // If the implicit attribute has its own declaration (e.g.,
+                        // `self.state: int`), that declaration takes precedence for
+                        // instance member lookup, even if the class has a declared type.
+                        if let PlaceAndQualifiers {
+                            place: Place::Defined(_, TypeOrigin::Declared, _, _),
+                            ..
+                        } = implicit_member.inner
                         {
+                            return implicit_member;
+                        }
+
+                        if let Some(implicit_ty) = implicit_member.ignore_possibly_undefined() {
                             if declaredness == Definedness::AlwaysDefined {
                                 // If a symbol is definitely declared, and we see
                                 // attribute assignments in methods of the class,

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -4760,15 +4760,43 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                             return false;
                         }
 
+                        let dunder_set_lookup = meta_attr_ty.class_member(db, "__set__".into());
                         let assignable_to_meta_attr =
-                            if let Place::Defined(meta_dunder_set, _, _, _) =
-                                meta_attr_ty.class_member(db, "__set__".into()).place
+                            if let Place::Defined(meta_dunder_set, _, dunder_set_boundness, _) =
+                                dunder_set_lookup.place
                             {
+                                // When `__set__` is only possibly defined (some union elements have it, some don't),
+                                // we need to check the `__set__` call only for elements that actually have `__set__`.
+                                // Otherwise, we'd be passing a union containing non-descriptor types as `self` to
+                                // `__set__`, which would fail. We also need to check normal assignment for non-descriptor
+                                // elements.
+                                let (descriptor_ty, non_descriptor_ty) = if dunder_set_boundness
+                                    == Definedness::PossiblyUndefined
+                                    && let Type::Union(union) = meta_attr_ty
+                                {
+                                    let with_set = union.filter(db, |elem| {
+                                        !elem
+                                            .class_member(db, "__set__".into())
+                                            .place
+                                            .is_undefined()
+                                    });
+                                    let without_set = union.filter(db, |elem| {
+                                        elem.class_member(db, "__set__".into()).place.is_undefined()
+                                    });
+                                    (with_set, Some(without_set))
+                                } else {
+                                    (meta_attr_ty, None)
+                                };
+
                                 // TODO: We could use the annotated parameter type of `__set__` as
                                 // type context here.
                                 let dunder_set_result = meta_dunder_set.try_call(
                                     db,
-                                    &CallArguments::positional([meta_attr_ty, object_ty, value_ty]),
+                                    &CallArguments::positional([
+                                        descriptor_ty,
+                                        object_ty,
+                                        value_ty,
+                                    ]),
                                 );
 
                                 if emit_diagnostics {
@@ -4783,7 +4811,36 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                     }
                                 }
 
-                                dunder_set_result.is_ok()
+                                let descriptor_ok = dunder_set_result.is_ok();
+
+                                // For union elements without `__set__`, the value
+                                // shadows the class attribute in the instance dict.
+                                // Check against instance member type if declared.
+                                let non_descriptor_ok =
+                                    if non_descriptor_ty.is_some_and(|ty| !ty.is_never()) {
+                                        if let PlaceAndQualifiers {
+                                            place: Place::Defined(instance_attr_ty, _, _, _),
+                                            qualifiers,
+                                        } = object_ty.instance_member(db, attribute)
+                                        {
+                                            let value_ty = infer_value_ty(
+                                                self,
+                                                TypeContext::new(Some(instance_attr_ty)),
+                                            );
+                                            if invalid_assignment_to_final(self, qualifiers) {
+                                                return false;
+                                            }
+                                            ensure_assignable_to(self, value_ty, instance_attr_ty)
+                                        } else {
+                                            // No instance member declared; value shadows
+                                            // the class attribute
+                                            true
+                                        }
+                                    } else {
+                                        true
+                                    };
+
+                                descriptor_ok && non_descriptor_ok
                             } else {
                                 let value_ty =
                                     infer_value_ty(self, TypeContext::new(Some(meta_attr_ty)));
@@ -4914,14 +4971,41 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                         }
 
                         let assignable_to_meta_attr =
-                            if let Place::Defined(meta_dunder_set, _, _, _) =
+                            if let Place::Defined(meta_dunder_set, _, dunder_set_boundness, _) =
                                 meta_attr_ty.class_member(db, "__set__".into()).place
                             {
+                                // When `__set__` is only possibly defined (some union elements have it, some don't),
+                                // we need to check the `__set__` call only for elements that actually have `__set__`.
+                                // Otherwise, we'd be passing a union containing non-descriptor types as `self` to
+                                // `__set__`, which would fail. We also need to check normal assignment for non-descriptor
+                                // elements.
+                                let (descriptor_ty, non_descriptor_ty) = if dunder_set_boundness
+                                    == Definedness::PossiblyUndefined
+                                    && let Type::Union(union) = meta_attr_ty
+                                {
+                                    let with_set = union.filter(db, |elem| {
+                                        !elem
+                                            .class_member(db, "__set__".into())
+                                            .place
+                                            .is_undefined()
+                                    });
+                                    let without_set = union.filter(db, |elem| {
+                                        elem.class_member(db, "__set__".into()).place.is_undefined()
+                                    });
+                                    (with_set, Some(without_set))
+                                } else {
+                                    (meta_attr_ty, None)
+                                };
+
                                 // TODO: We could use the annotated parameter type of `__set__` as
                                 // type context here.
                                 let dunder_set_result = meta_dunder_set.try_call(
                                     db,
-                                    &CallArguments::positional([meta_attr_ty, object_ty, value_ty]),
+                                    &CallArguments::positional([
+                                        descriptor_ty,
+                                        object_ty,
+                                        value_ty,
+                                    ]),
                                 );
 
                                 if emit_diagnostics {
@@ -4936,7 +5020,27 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                     }
                                 }
 
-                                dunder_set_result.is_ok()
+                                let descriptor_ok = dunder_set_result.is_ok();
+
+                                // For union elements without `__set__`, the value
+                                // is written directly to the class dict. Check that
+                                // the value is assignable to the type of those elements.
+                                let non_descriptor_ok = if let Some(non_desc_ty) = non_descriptor_ty
+                                {
+                                    if non_desc_ty.is_never() {
+                                        true
+                                    } else {
+                                        let value_ty = infer_value_ty(
+                                            self,
+                                            TypeContext::new(Some(non_desc_ty)),
+                                        );
+                                        ensure_assignable_to(self, value_ty, non_desc_ty)
+                                    }
+                                } else {
+                                    true
+                                };
+
+                                descriptor_ok && non_descriptor_ok
                             } else {
                                 let value_ty =
                                     infer_value_ty(self, TypeContext::new(Some(meta_attr_ty)));