[ty] Skip eagerly evaluated scopes for attribute storing (#20856)

## Summary

Fix https://github.com/astral-sh/ty/issues/664

This PR adds support for storing attributes in comprehension scopes (any
eager scope.)

For example in the following code we infer type of `z` correctly:

```py
class C:
    def __init__(self):
        [None for self.z in range(1)]
reveal_type(C().z) # previously [unresolved-attribute] but now shows Unknown | int
```

The fix works by adjusting the following logics:

To identify if an attriute is an assignment to self or cls we need to
check the scope is a method. To allow comprehension scopes here we skip
any eager scope in the check.
Also at this stage the code checks if self or the first method argument
is shadowed by another binding that eager scope to prevent this:

```py
class D:
    g: int

class C:
    def __init__(self):
        [[None for self.g in range(1)] for self in [D()]]
reveal_type(C().g) # [unresolved-attribute]
```

When determining scopes that attributes might be defined after
collecting all the methods of the class the code also returns any
decendant scope that is eager and only has eager parents until the
method scope.

When checking reachability of a attribute definition if the attribute is
defined in an eager scope we use the reachability of the first non eager
scope which must be a method. This allows attributes to be marked as
reachable and be seen.


There are also which I didn't add support for:

```py
class C:
    def __init__(self):
        def f():
            [None for self.z in range(1)]
        f()

reveal_type(C().z) # [unresolved-attribute]
```

In the above example we will not even return the comprehension scope as
an attribute scope because there is a non eager scope (`f` function)
between the comprehension and the `__init__` method

---------

Co-authored-by: Carl Meyer <carl@astral.sh>

crates/ty_python_semantic/resources/mdtest/attributes.md

@@ -369,6 +369,11 @@ reveal_type(c_instance.y)  # revealed: Unknown | int
 
 #### Attributes defined in comprehensions
 
+```toml
+[environment]
+python-version = "3.12"
+```
+
 ```py
 class TupleIterator:
     def __next__(self) -> tuple[int, str]:
@@ -380,19 +385,9 @@ class TupleIterable:
 
 class C:
     def __init__(self) -> None:
-        # TODO: Should not emit this diagnostic
-        # error: [unresolved-attribute]
         [... for self.a in range(3)]
-        # TODO: Should not emit this diagnostic
-        # error: [unresolved-attribute]
-        # error: [unresolved-attribute]
         [... for (self.b, self.c) in TupleIterable()]
-        # TODO: Should not emit this diagnostic
-        # error: [unresolved-attribute]
-        # error: [unresolved-attribute]
         [... for self.d in range(3) for self.e in range(3)]
-        # TODO: Should not emit this diagnostic
-        # error: [unresolved-attribute]
         [[... for self.f in range(3)] for _ in range(3)]
         [[... for self.g in range(3)] for self in [D()]]
 
@@ -401,35 +396,74 @@ class D:
 
 c_instance = C()
 
-# TODO: no error, reveal Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.a)  # revealed: Unknown
+reveal_type(c_instance.a)  # revealed: Unknown | int
 
-# TODO: no error, reveal Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.b)  # revealed: Unknown
+reveal_type(c_instance.b)  # revealed: Unknown | int
 
-# TODO: no error, reveal Unknown | str
-# error: [unresolved-attribute]
-reveal_type(c_instance.c)  # revealed: Unknown
+reveal_type(c_instance.c)  # revealed: Unknown | str
 
-# TODO: no error, reveal Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.d)  # revealed: Unknown
+reveal_type(c_instance.d)  # revealed: Unknown | int
 
-# TODO: no error, reveal Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.e)  # revealed: Unknown
+reveal_type(c_instance.e)  # revealed: Unknown | int
 
-# TODO: no error, reveal Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.f)  # revealed: Unknown
+reveal_type(c_instance.f)  # revealed: Unknown | int
 
 # This one is correctly not resolved as an attribute:
 # error: [unresolved-attribute]
 reveal_type(c_instance.g)  # revealed: Unknown
 ```
 
+It does not matter how much the comprehension is nested.
+
+Similarly attributes defined by the comprehension in a generic method are recognized.
+
+```py
+class C:
+    def f[T](self):
+        [... for self.a in [1]]
+        [[... for self.b in [1]] for _ in [1]]
+
+c_instance = C()
+
+reveal_type(c_instance.a)  # revealed: Unknown | int
+reveal_type(c_instance.b)  # revealed: Unknown | int
+```
+
+If the comprehension is inside another scope like function then that attribute is not inferred.
+
+```py
+class C:
+    def __init__(self):
+        def f():
+            # error: [unresolved-attribute]
+            [... for self.a in [1]]
+
+        def g():
+            # error: [unresolved-attribute]
+            [... for self.b in [1]]
+        g()
+
+c_instance = C()
+
+# This attribute is in the function f and is not reachable
+# error: [unresolved-attribute]
+reveal_type(c_instance.a)  # revealed: Unknown
+
+# error: [unresolved-attribute]
+reveal_type(c_instance.b)  # revealed: Unknown
+```
+
+If the comprehension is nested in any other eager scope it still can assign attributes.
+
+```py
+class C:
+    def __init__(self):
+        class D:
+            [[... for self.a in [1]] for _ in [1]]
+
+reveal_type(C().a)  # revealed: Unknown | int
+```
+
 #### Conditionally declared / bound attributes
 
 We currently treat implicit instance attributes to be bound, even if they are only conditionally

crates/ty_python_semantic/src/semantic_index.rs

@@ -1,4 +1,4 @@
-use std::iter::FusedIterator;
+use std::iter::{FusedIterator, once};
 use std::sync::Arc;
 
 use ruff_db::files::File;
@@ -148,29 +148,56 @@ pub(crate) fn attribute_declarations<'db, 's>(
 ///
 /// Only call this when doing type inference on the same file as `class_body_scope`, otherwise it
 /// introduces a direct dependency on that file's AST.
-pub(crate) fn attribute_scopes<'db, 's>(
+pub(crate) fn attribute_scopes<'db>(
     db: &'db dyn Db,
     class_body_scope: ScopeId<'db>,
-) -> impl Iterator<Item = FileScopeId> + use<'s, 'db> {
+) -> impl Iterator<Item = FileScopeId> + 'db {
     let file = class_body_scope.file(db);
     let index = semantic_index(db, file);
     let class_scope_id = class_body_scope.file_scope_id(db);
+    ChildrenIter::new(&index.scopes, class_scope_id)
+        .filter_map(move |(child_scope_id, scope)| {
+            let (function_scope_id, function_scope) =
+                if scope.node().scope_kind() == ScopeKind::TypeParams {
+                    // This could be a generic method with a type-params scope.
+                    // Go one level deeper to find the function scope. The first
+                    // descendant is the (potential) function scope.
+                    let function_scope_id = scope.descendants().start;
+                    (function_scope_id, index.scope(function_scope_id))
+                } else {
+                    (child_scope_id, scope)
+                };
+            function_scope.node().as_function()?;
+            Some(function_scope_id)
+        })
+        .flat_map(move |func_id| {
+            // Add any descendent scope that is eager and have eager scopes between the scope
+            // and the method scope. Since attributes can be defined in this scope.
+            let nested = index.descendent_scopes(func_id).filter_map(move |(id, s)| {
+                let is_eager = s.kind().is_eager();
+                let parents_are_eager = {
+                    let mut all_parents_eager = true;
+                    let mut current = Some(id);
 
-    ChildrenIter::new(&index.scopes, class_scope_id).filter_map(move |(child_scope_id, scope)| {
-        let (function_scope_id, function_scope) =
-            if scope.node().scope_kind() == ScopeKind::TypeParams {
-                // This could be a generic method with a type-params scope.
-                // Go one level deeper to find the function scope. The first
-                // descendant is the (potential) function scope.
-                let function_scope_id = scope.descendants().start;
-                (function_scope_id, index.scope(function_scope_id))
-            } else {
-                (child_scope_id, scope)
-            };
+                    while let Some(scope_id) = current {
+                        if scope_id == func_id {
+                            break;
+                        }
+                        let scope = index.scope(scope_id);
+                        if !scope.is_eager() {
+                            all_parents_eager = false;
+                            break;
+                        }
+                        current = scope.parent();
+                    }
 
-        function_scope.node().as_function()?;
-        Some(function_scope_id)
-    })
+                    all_parents_eager
+                };
+
+                (parents_are_eager && is_eager).then_some(id)
+            });
+            once(func_id).chain(nested)
+        })
 }
 
 /// Returns the module global scope of `file`.

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -186,29 +186,34 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         self.current_scope_info().file_scope_id
     }
 
-    /// Returns the scope ID of the surrounding class body scope if the current scope
-    /// is a method inside a class body. Returns `None` otherwise, e.g. if the current
-    /// scope is a function body outside of a class, or if the current scope is not a
+    /// Returns the scope ID of the current scope if the current scope
+    /// is a method inside a class body or an eagerly executed scope inside a method.
+    /// Returns `None` otherwise, e.g. if the current scope is a function body outside of a class, or if the current scope is not a
     /// function body.
-    fn is_method_of_class(&self) -> Option<FileScopeId> {
-        let mut scopes_rev = self.scope_stack.iter().rev();
+    fn is_method_or_eagerly_executed_in_method(&self) -> Option<FileScopeId> {
+        let mut scopes_rev = self
+            .scope_stack
+            .iter()
+            .rev()
+            .skip_while(|scope| self.scopes[scope.file_scope_id].is_eager());
         let current = scopes_rev.next()?;
 
         if self.scopes[current.file_scope_id].kind() != ScopeKind::Function {
             return None;
         }
 
+        let maybe_method = current.file_scope_id;
         let parent = scopes_rev.next()?;
 
         match self.scopes[parent.file_scope_id].kind() {
-            ScopeKind::Class => Some(parent.file_scope_id),
+            ScopeKind::Class => Some(maybe_method),
             ScopeKind::TypeParams => {
                 // If the function is generic, the parent scope is an annotation scope.
                 // In this case, we need to go up one level higher to find the class scope.
                 let grandparent = scopes_rev.next()?;
 
                 if self.scopes[grandparent.file_scope_id].kind() == ScopeKind::Class {
-                    Some(grandparent.file_scope_id)
+                    Some(maybe_method)
                 } else {
                     None
                 }
@@ -217,6 +222,32 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         }
     }
 
+    /// Checks if a symbol name is bound in any intermediate eager scopes
+    /// between the current scope and the specified method scope.
+    ///
+    fn is_symbol_bound_in_intermediate_eager_scopes(
+        &self,
+        symbol_name: &str,
+        method_scope_id: FileScopeId,
+    ) -> bool {
+        for scope_info in self.scope_stack.iter().rev() {
+            let scope_id = scope_info.file_scope_id;
+
+            if scope_id == method_scope_id {
+                break;
+            }
+
+            if let Some(symbol_id) = self.place_tables[scope_id].symbol_id(symbol_name) {
+                let symbol = self.place_tables[scope_id].symbol(symbol_id);
+                if symbol.is_bound() {
+                    return true;
+                }
+            }
+        }
+
+        false
+    }
+
     /// Push a new loop, returning the outer loop, if any.
     fn push_loop(&mut self) -> Option<Loop> {
         self.current_scope_info_mut()
@@ -1700,7 +1731,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 self.visit_expr(&node.annotation);
                 if let Some(value) = &node.value {
                     self.visit_expr(value);
-                    if self.is_method_of_class().is_some() {
+                    if self.is_method_or_eagerly_executed_in_method().is_some() {
                         // Record the right-hand side of the assignment as a standalone expression
                         // if we're inside a method. This allows type inference to infer the type
                         // of the value for annotated assignments like `self.CONSTANT: Final = 1`,
@@ -2372,14 +2403,21 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
             | ast::Expr::Attribute(ast::ExprAttribute { ctx, .. })
             | ast::Expr::Subscript(ast::ExprSubscript { ctx, .. }) => {
                 if let Some(mut place_expr) = PlaceExpr::try_from_expr(expr) {
-                    if self.is_method_of_class().is_some() {
+                    if let Some(method_scope_id) = self.is_method_or_eagerly_executed_in_method() {
                         if let PlaceExpr::Member(member) = &mut place_expr {
                             if member.is_instance_attribute_candidate() {
                                 // We specifically mark attribute assignments to the first parameter of a method,
                                 // i.e. typically `self` or `cls`.
-                                let accessed_object_refers_to_first_parameter = self
-                                    .current_first_parameter_name
-                                    .is_some_and(|first| member.symbol_name() == first);
+                                // However, we must check that the symbol hasn't been shadowed by an intermediate
+                                // scope (e.g., a comprehension variable: `for self in [...]`).
+                                let accessed_object_refers_to_first_parameter =
+                                    self.current_first_parameter_name.is_some_and(|first| {
+                                        member.symbol_name() == first
+                                            && !self.is_symbol_bound_in_intermediate_eager_scopes(
+                                                first,
+                                                method_scope_id,
+                                            )
+                                    });
 
                                 if accessed_object_refers_to_first_parameter {
                                     member.mark_instance_attribute();

crates/ty_python_semantic/src/types/class.rs

@@ -3119,30 +3119,47 @@ impl<'db> ClassLiteral<'db> {
             union_of_inferred_types = union_of_inferred_types.add(Type::unknown());
         }
 
-        for (attribute_assignments, method_scope_id) in
+        for (attribute_assignments, attribute_binding_scope_id) in
             attribute_assignments(db, class_body_scope, &name)
         {
-            let method_scope = index.scope(method_scope_id);
-            if !is_valid_scope(method_scope) {
+            let binding_scope = index.scope(attribute_binding_scope_id);
+            if !is_valid_scope(binding_scope) {
                 continue;
             }
 
-            // The attribute assignment inherits the reachability of the method which contains it
-            let is_method_reachable = if let Some(method_def) = method_scope.node().as_function() {
-                let method = index.expect_single_definition(method_def);
-                let method_place = class_table
-                    .symbol_id(&method_def.node(&module).name)
-                    .unwrap();
-                class_map
-                    .all_reachable_symbol_bindings(method_place)
-                    .find_map(|bind| {
-                        (bind.binding.is_defined_and(|def| def == method))
-                            .then(|| class_map.binding_reachability(db, &bind))
-                    })
-                    .unwrap_or(Truthiness::AlwaysFalse)
-            } else {
-                Truthiness::AlwaysFalse
+            let scope_for_reachability_analysis = {
+                if binding_scope.node().as_function().is_some() {
+                    binding_scope
+                } else if binding_scope.is_eager() {
+                    let mut eager_scope_parent = binding_scope;
+                    while eager_scope_parent.is_eager()
+                        && let Some(parent) = eager_scope_parent.parent()
+                    {
+                        eager_scope_parent = index.scope(parent);
+                    }
+                    eager_scope_parent
+                } else {
+                    binding_scope
+                }
             };
+
+            // The attribute assignment inherits the reachability of the method which contains it
+            let is_method_reachable =
+                if let Some(method_def) = scope_for_reachability_analysis.node().as_function() {
+                    let method = index.expect_single_definition(method_def);
+                    let method_place = class_table
+                        .symbol_id(&method_def.node(&module).name)
+                        .unwrap();
+                    class_map
+                        .all_reachable_symbol_bindings(method_place)
+                        .find_map(|bind| {
+                            (bind.binding.is_defined_and(|def| def == method))
+                                .then(|| class_map.binding_reachability(db, &bind))
+                        })
+                        .unwrap_or(Truthiness::AlwaysFalse)
+                } else {
+                    Truthiness::AlwaysFalse
+                };
             if is_method_reachable.is_always_false() {
                 continue;
             }