Merge branch 'comprehension-assign-attr-unresolved' into david/type-of-self-in-methods-integration-2

crates/ty_python_semantic/resources/mdtest/attributes.md

@@ -377,6 +377,11 @@ reveal_type(c_instance.y)  # revealed: Unknown | int
 
 #### Attributes defined in comprehensions
 
+```toml
+[environment]
+python-version = "3.12"
+```
+
 ```py
 class IntIterator:
     def __next__(self) -> int:
@@ -417,35 +422,73 @@ class D:
 
 c_instance = C()
 
-# TODO: no error, reveal Unknown | int
+reveal_type(c_instance.a)  # revealed: Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.a)  # revealed: Unknown
 
-# TODO: no error, reveal Unknown | int
+reveal_type(c_instance.b)  # revealed: Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.b)  # revealed: Unknown
 
-# TODO: no error, reveal Unknown | str
+reveal_type(c_instance.c)  # revealed: Unknown | str
-# error: [unresolved-attribute]
-reveal_type(c_instance.c)  # revealed: Unknown
 
-# TODO: no error, reveal Unknown | int
+reveal_type(c_instance.d)  # revealed: Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.d)  # revealed: Unknown
 
-# TODO: no error, reveal Unknown | int
+reveal_type(c_instance.e)  # revealed: Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.e)  # revealed: Unknown
 
-# TODO: no error, reveal Unknown | int
+reveal_type(c_instance.f)  # revealed: Unknown | int
-# error: [unresolved-attribute]
-reveal_type(c_instance.f)  # revealed: Unknown
 
 # 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():
+            [... for self.a in IntIterable()]
+
+        def g():
+            [... for self.b in IntIterable()]
+        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
+
+# TODO: Even though g method is called and is reachable we do not record this attribute assignment
+# 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 IntIterable()] for _ in IntIterable()]
+
+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;
@@ -154,29 +154,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)| {
+                    while let Some(scope_id) = current {
-        let (function_scope_id, function_scope) =
+                        if scope_id == func_id {
-            if scope.node().scope_kind() == ScopeKind::TypeParams {
+                            break;
-                // This could be a generic method with a type-params scope.
+                        }
-                // Go one level deeper to find the function scope. The first
+                        let scope = index.scope(scope_id);
-                // descendant is the (potential) function scope.
+                        if !scope.is_eager() {
-                let function_scope_id = scope.descendants().start;
+                            all_parents_eager = false;
-                (function_scope_id, index.scope(function_scope_id))
+                            break;
-            } else {
+                        }
-                (child_scope_id, scope)
+                        current = scope.parent();
-            };
+                    }
 
-        function_scope.node().as_function()?;
+                    all_parents_eager
-        Some(function_scope_id)
+                };
-    })
+
+                (parents_are_eager && is_eager).then_some(id)
+            });
+            once(func_id).chain(nested)
+        })
 }
 
 /// Returns the module global scope of `file`.
@@ -751,7 +778,10 @@ mod tests {
     use crate::semantic_index::scope::{FileScopeId, Scope, ScopeKind};
     use crate::semantic_index::symbol::ScopedSymbolId;
     use crate::semantic_index::use_def::UseDefMap;
-    use crate::semantic_index::{global_scope, place_table, semantic_index, use_def_map};
+    use crate::semantic_index::{
+        attribute_declarations, attribute_scopes, global_scope, place_table, semantic_index,
+        use_def_map,
+    };
 
     impl UseDefMap<'_> {
         fn first_public_binding(&self, symbol: ScopedSymbolId) -> Option<Definition<'_>> {
@@ -1128,6 +1158,124 @@ def f(a: str, /, b: str, c: int = 1, *args, d: int = 2, **kwargs):
         }
     }
 
+    /// Test case to validate that comprehensions inside a method are correctly marked as attribute
+    /// scopes
+    #[test]
+    fn comprehension_in_method_instance_attribute() {
+        let TestCase { db, file } = test_case(
+            "
+class C:
+    def __init__(self):
+        [None for self.z in range(1)]
+",
+        );
+
+        let index = semantic_index(&db, file);
+
+        let [(class_scope_id, class_scope)] = index
+            .child_scopes(FileScopeId::global())
+            .collect::<Vec<_>>()[..]
+        else {
+            panic!("expected one child scope (the class)")
+        };
+
+        assert_eq!(class_scope.kind(), ScopeKind::Class);
+        let class_scope = class_scope_id.to_scope_id(&db, file);
+
+        let method_scopes: Vec<_> = index.child_scopes(class_scope_id).collect();
+        assert_eq!(method_scopes.len(), 1, "expected __init__");
+        let (method_scope_id, method_scope) = method_scopes[0];
+        assert_eq!(method_scope.kind(), ScopeKind::Function);
+
+        let comp_scopes: Vec<_> = index.child_scopes(method_scope_id).collect();
+        assert_eq!(comp_scopes.len(), 1, "expected one comprehension scope");
+        let (comprehension_scope_id, comprehension_scope) = comp_scopes[0];
+        assert_eq!(comprehension_scope.kind(), ScopeKind::Comprehension);
+
+        let attr_scopes: Vec<_> = attribute_scopes(&db, class_scope).collect();
+        assert!(
+            attr_scopes.contains(&method_scope_id),
+            "attribute_scopes should include the method scope"
+        );
+        assert!(
+            attr_scopes.contains(&comprehension_scope_id),
+            "attribute_scopes should include the comprehension scope"
+        );
+
+        let comprehension_place_table = index.place_table(comprehension_scope_id);
+        let members: Vec<_> = comprehension_place_table.members().collect();
+
+        let has_z_instance_attr = members
+            .iter()
+            .any(|member| member.as_instance_attribute() == Some("z"));
+        assert!(
+            has_z_instance_attr,
+            "self.z should be marked as an instance attribute in the comprehension scope found members: {members:?}"
+        );
+
+        let z_declarations: Vec<_> = attribute_declarations(&db, class_scope, "z")
+            .map(|(decls, scope_id)| {
+                let decl_vec: Vec<_> = decls.collect();
+                (decl_vec, scope_id)
+            })
+            .collect();
+        assert!(
+            !z_declarations.is_empty(),
+            "attribute_declarations should find declarations for z"
+        );
+
+        // Verify that at least one declaration of z is in the comprehension scope
+        let has_comp_declaration = z_declarations
+            .iter()
+            .any(|(_, scope_id)| *scope_id == comprehension_scope_id);
+        assert!(
+            has_comp_declaration,
+            "attribute_declarations should include a declaration from the comprehension scope"
+        );
+    }
+
+    /// Test case to validate that when first method argument is shadowed by a comprehension variable,
+    /// attributes accessed on the shadowed argument should NOT be tracked as instance attributes
+    /// of the containing class.
+    #[test]
+    fn comprehension_shadowing_self() {
+        let TestCase { db, file } = test_case(
+            "
+class D:
+    g: int
+
+class C:
+    def __init__(self):
+        [[None for self.g in range(1)] for self in [D()]]
+",
+        );
+
+        let index = semantic_index(&db, file);
+
+        let all_scopes: Vec<_> = index.child_scopes(FileScopeId::global()).collect();
+
+        let (class_c_scope_id, class_c_scope) = all_scopes[1];
+        assert_eq!(class_c_scope.kind(), ScopeKind::Class);
+        let class_c_scope = class_c_scope_id.to_scope_id(&db, file);
+
+        let attr_scopes: Vec<_> = attribute_scopes(&db, class_c_scope).collect();
+        let mut has_g_attribute = false;
+        for scope_id in &attr_scopes {
+            let place_table = index.place_table(*scope_id);
+            if place_table
+                .member_id_by_instance_attribute_name("g")
+                .is_some()
+            {
+                has_g_attribute = true;
+                break;
+            }
+        }
+        assert!(
+            !has_g_attribute,
+            "self.g should NOT be tracked as an attribute of C because 'self' is shadowed by the outer comprehension variable"
+        );
+    }
+
     /// Test case to validate that the `x` variable used in the comprehension is referencing the
     /// `x` variable defined by the inner generator (`for x in iter2`) and not the outer one.
     #[test]

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -184,29 +184,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
+    /// Returns the scope ID of the method scope if the current scope
-    /// is a method inside a class body. Returns `None` otherwise, e.g. if the current
+    /// is a method inside a class body or current scope is in eagerly executed scope in a method.
-    /// scope is a function body outside of a class, or if the current scope is not a
+    /// 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> {
+    fn is_method_or_eagerly_executed_in_method(&self) -> Option<FileScopeId> {
-        let mut scopes_rev = self.scope_stack.iter().rev();
+        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
                 }
@@ -215,6 +220,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()
@@ -1647,7 +1678,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`,
@@ -2319,14 +2350,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
+                                // However, we must check that the symbol hasn't been shadowed by an intermediate
-                                    .current_first_parameter_name
+                                // scope (e.g., a comprehension variable: `for self in [...]`).
-                                    .is_some_and(|first| member.symbol_name() == first);
+                                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

@@ -3227,30 +3227,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);
+            let binding_scope = index.scope(attribute_binding_scope_id);
-            if !is_valid_scope(method_scope) {
+            if !is_valid_scope(binding_scope) {
                 continue;
             }
 
-            // The attribute assignment inherits the reachability of the method which contains it
+            let scope_for_reachability_analysis = {
-            let is_method_reachable = if let Some(method_def) = method_scope.node().as_function() {
+                if binding_scope.node().as_function().is_some() {
-                let method = index.expect_single_definition(method_def);
+                    binding_scope
-                let method_place = class_table
+                } else if binding_scope.is_eager() {
-                    .symbol_id(&method_def.node(&module).name)
+                    let mut eager_scope_parent = binding_scope;
-                    .unwrap();
+                    while eager_scope_parent.is_eager()
-                class_map
+                        && let Some(parent) = eager_scope_parent.parent()
-                    .all_reachable_symbol_bindings(method_place)
+                    {
-                    .find_map(|bind| {
+                        eager_scope_parent = index.scope(parent);
-                        (bind.binding.is_defined_and(|def| def == method))
+                    }
-                            .then(|| class_map.binding_reachability(db, &bind))
+                    eager_scope_parent
-                    })
+                } else {
-                    .unwrap_or(Truthiness::AlwaysFalse)
+                    binding_scope
-            } else {
+                }
-                Truthiness::AlwaysFalse
             };
+
+            // 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;
             }