Skip to content

Commit 3cf44e4

Browse files
authored
[red-knot] Implicit instance attributes in generic methods (#17769)
## Summary Add the ability to detect instance attribute assignments in class methods that are generic. This does not address the code duplication mentioned in #16928. I can open a ticket for this after this has been merged. closes #16928 ## Test Plan Added regression test.
1 parent 17050e2 commit 3cf44e4

File tree

4 files changed

+69
-11
lines changed

4 files changed

+69
-11
lines changed
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# Regression test for an issue that came up while working
2+
# on https://github.com/astral-sh/ruff/pull/17769
3+
4+
class C:
5+
def method[T](self, x: T) -> T:
6+
def inner():
7+
self.attr = 1
8+
9+
C().attr

crates/red_knot_python_semantic/resources/mdtest/attributes.md

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1949,6 +1949,27 @@ reveal_type(C.a_type) # revealed: type
19491949
reveal_type(C.a_none) # revealed: None
19501950
```
19511951

1952+
### Generic methods
1953+
1954+
We also detect implicit instance attributes on methods that are themselves generic. We have an extra
1955+
test for this because generic functions have an extra type-params scope in between the function body
1956+
scope and the outer scope, so we need to make sure that our implementation can still recognize `f`
1957+
as a method of `C` here:
1958+
1959+
```toml
1960+
[environment]
1961+
python-version = "3.12"
1962+
```
1963+
1964+
```py
1965+
class C:
1966+
def f[T](self, t: T) -> T:
1967+
self.x: int = 1
1968+
return t
1969+
1970+
reveal_type(C().x) # revealed: int
1971+
```
1972+
19521973
## Enum classes
19531974

19541975
Enums are not supported yet; attribute access on an enum class is inferred as `Todo`.

crates/red_knot_python_semantic/src/semantic_index.rs

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@ use crate::semantic_index::builder::SemanticIndexBuilder;
1818
use crate::semantic_index::definition::{Definition, DefinitionNodeKey, Definitions};
1919
use crate::semantic_index::expression::Expression;
2020
use crate::semantic_index::symbol::{
21-
FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopedSymbolId, SymbolTable,
21+
FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopeKind, ScopedSymbolId,
22+
SymbolTable,
2223
};
2324
use crate::semantic_index::use_def::{EagerBindingsKey, ScopedEagerBindingsId, UseDefMap};
2425
use crate::Db;
@@ -109,12 +110,26 @@ pub(crate) fn attribute_assignments<'db, 's>(
109110
let index = semantic_index(db, file);
110111
let class_scope_id = class_body_scope.file_scope_id(db);
111112

112-
ChildrenIter::new(index, class_scope_id).filter_map(|(file_scope_id, maybe_method)| {
113-
maybe_method.node().as_function()?;
114-
let attribute_table = index.instance_attribute_table(file_scope_id);
113+
ChildrenIter::new(index, class_scope_id).filter_map(|(child_scope_id, scope)| {
114+
let (function_scope_id, function_scope) =
115+
if scope.node().scope_kind() == ScopeKind::Annotation {
116+
// This could be a generic method with a type-params scope.
117+
// Go one level deeper to find the function scope. The first
118+
// descendant is the (potential) function scope.
119+
let function_scope_id = scope.descendants().start;
120+
(function_scope_id, index.scope(function_scope_id))
121+
} else {
122+
(child_scope_id, scope)
123+
};
124+
125+
function_scope.node().as_function()?;
126+
let attribute_table = index.instance_attribute_table(function_scope_id);
115127
let symbol = attribute_table.symbol_id_by_name(name)?;
116-
let use_def = &index.use_def_maps[file_scope_id];
117-
Some((use_def.instance_attribute_bindings(symbol), file_scope_id))
128+
let use_def = &index.use_def_maps[function_scope_id];
129+
Some((
130+
use_def.instance_attribute_bindings(symbol),
131+
function_scope_id,
132+
))
118133
})
119134
}
120135

crates/red_knot_python_semantic/src/semantic_index/builder.rs

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -183,13 +183,26 @@ impl<'db> SemanticIndexBuilder<'db> {
183183
fn is_method_of_class(&self) -> Option<FileScopeId> {
184184
let mut scopes_rev = self.scope_stack.iter().rev();
185185
let current = scopes_rev.next()?;
186+
187+
if self.scopes[current.file_scope_id].kind() != ScopeKind::Function {
188+
return None;
189+
}
190+
186191
let parent = scopes_rev.next()?;
187192

188-
match (
189-
self.scopes[current.file_scope_id].kind(),
190-
self.scopes[parent.file_scope_id].kind(),
191-
) {
192-
(ScopeKind::Function, ScopeKind::Class) => Some(parent.file_scope_id),
193+
match self.scopes[parent.file_scope_id].kind() {
194+
ScopeKind::Class => Some(parent.file_scope_id),
195+
ScopeKind::Annotation => {
196+
// If the function is generic, the parent scope is an annotation scope.
197+
// In this case, we need to go up one level higher to find the class scope.
198+
let grandparent = scopes_rev.next()?;
199+
200+
if self.scopes[grandparent.file_scope_id].kind() == ScopeKind::Class {
201+
Some(grandparent.file_scope_id)
202+
} else {
203+
None
204+
}
205+
}
193206
_ => None,
194207
}
195208
}

0 commit comments

Comments
 (0)