[red-knot] Implicit instance attributes (#15811)

## Summary

Add support for implicitly-defined instance attributes, i.e. support
type inference for cases like this:
```py
class C:
    def __init__(self) -> None:
        self.x: int = 1
        self.y = None

reveal_type(C().x)  # int
reveal_type(C().y)  # Unknown | None
```

## Benchmarks

Codspeed reports no change in a cold-cache benchmark, and a -1%
regression in the incremental benchmark. On `black`'s `src` folder, I
don't see a statistically significant difference between the branches:

| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|:---|---:|---:|---:|---:|
| `./red_knot_main check --project /home/shark/black/src` | 133.7 ± 9.5 | 126.7 | 164.7 | 1.01 ± 0.08 |
| `./red_knot_feature check --project /home/shark/black/src` | 132.2 ± 5.1 | 118.1 | 140.9 | 1.00 |

## Test Plan

Updated and new Markdown tests
This commit is contained in:
David Peter 2025-02-03 19:34:23 +01:00 committed by GitHub
parent dc5e922221
commit 102c2eec12
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 657 additions and 53 deletions

View file

@ -14,14 +14,15 @@ use crate::ast_node_ref::AstNodeRef;
use crate::module_name::ModuleName;
use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
use crate::semantic_index::ast_ids::AstIdsBuilder;
use crate::semantic_index::attribute_assignment::{AttributeAssignment, AttributeAssignments};
use crate::semantic_index::constraint::PatternConstraintKind;
use crate::semantic_index::definition::{
AssignmentDefinitionNodeRef, ComprehensionDefinitionNodeRef, Definition, DefinitionNodeKey,
DefinitionNodeRef, ForStmtDefinitionNodeRef, ImportFromDefinitionNodeRef,
};
use crate::semantic_index::expression::Expression;
use crate::semantic_index::expression::{Expression, ExpressionKind};
use crate::semantic_index::symbol::{
FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopedSymbolId,
FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopeKind, ScopedSymbolId,
SymbolTableBuilder,
};
use crate::semantic_index::use_def::{
@ -53,17 +54,24 @@ impl LoopState {
}
}
struct ScopeInfo {
file_scope_id: FileScopeId,
loop_state: LoopState,
}
pub(super) struct SemanticIndexBuilder<'db> {
// Builder state
db: &'db dyn Db,
file: File,
module: &'db ParsedModule,
scope_stack: Vec<(FileScopeId, LoopState)>,
scope_stack: Vec<ScopeInfo>,
/// The assignments we're currently visiting, with
/// the most recent visit at the end of the Vec
current_assignments: Vec<CurrentAssignment<'db>>,
/// The match case we're currently visiting.
current_match_case: Option<CurrentMatchCase<'db>>,
/// The name of the first function parameter of the innermost function that we're currently visiting.
current_first_parameter_name: Option<&'db str>,
/// Flow states at each `break` in the current loop.
loop_break_states: Vec<FlowSnapshot>,
@ -84,6 +92,7 @@ pub(super) struct SemanticIndexBuilder<'db> {
definitions_by_node: FxHashMap<DefinitionNodeKey, Definition<'db>>,
expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
imported_modules: FxHashSet<ModuleName>,
attribute_assignments: FxHashMap<FileScopeId, AttributeAssignments<'db>>,
}
impl<'db> SemanticIndexBuilder<'db> {
@ -95,6 +104,7 @@ impl<'db> SemanticIndexBuilder<'db> {
scope_stack: Vec::new(),
current_assignments: vec![],
current_match_case: None,
current_first_parameter_name: None,
loop_break_states: vec![],
try_node_context_stack_manager: TryNodeContextStackManager::default(),
@ -112,6 +122,8 @@ impl<'db> SemanticIndexBuilder<'db> {
expressions_by_node: FxHashMap::default(),
imported_modules: FxHashSet::default(),
attribute_assignments: FxHashMap::default(),
};
builder.push_scope_with_parent(NodeWithScopeRef::Module, None);
@ -123,7 +135,7 @@ impl<'db> SemanticIndexBuilder<'db> {
*self
.scope_stack
.last()
.map(|(scope, _)| scope)
.map(|ScopeInfo { file_scope_id, .. }| file_scope_id)
.expect("Always to have a root scope")
}
@ -131,14 +143,32 @@ impl<'db> SemanticIndexBuilder<'db> {
self.scope_stack
.last()
.expect("Always to have a root scope")
.1
.loop_state
}
/// 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
/// function body.
fn is_method_of_class(&self) -> Option<FileScopeId> {
let mut scopes_rev = self.scope_stack.iter().rev();
let current = scopes_rev.next()?;
let parent = scopes_rev.next()?;
match (
self.scopes[current.file_scope_id].kind(),
self.scopes[parent.file_scope_id].kind(),
) {
(ScopeKind::Function, ScopeKind::Class) => Some(parent.file_scope_id),
_ => None,
}
}
fn set_inside_loop(&mut self, state: LoopState) {
self.scope_stack
.last_mut()
.expect("Always to have a root scope")
.1 = state;
.loop_state = state;
}
fn push_scope(&mut self, node: NodeWithScopeRef) {
@ -171,16 +201,20 @@ impl<'db> SemanticIndexBuilder<'db> {
debug_assert_eq!(ast_id_scope, file_scope_id);
self.scope_stack.push((file_scope_id, LoopState::NotInLoop));
self.scope_stack.push(ScopeInfo {
file_scope_id,
loop_state: LoopState::NotInLoop,
});
}
fn pop_scope(&mut self) -> FileScopeId {
let (id, _) = self.scope_stack.pop().expect("Root scope to be present");
let ScopeInfo { file_scope_id, .. } =
self.scope_stack.pop().expect("Root scope to be present");
let children_end = self.scopes.next_index();
let scope = &mut self.scopes[id];
let scope = &mut self.scopes[file_scope_id];
scope.descendents = scope.descendents.start..children_end;
self.try_node_context_stack_manager.exit_scope();
id
file_scope_id
}
fn current_symbol_table(&mut self) -> &mut SymbolTableBuilder {
@ -404,6 +438,32 @@ impl<'db> SemanticIndexBuilder<'db> {
self.current_assignments.last_mut()
}
/// Records the fact that we saw an attribute assignment of the form
/// `object.attr: <annotation>( = …)` or `object.attr = <value>`.
fn register_attribute_assignment(
&mut self,
object: &ast::Expr,
attr: &'db ast::Identifier,
attribute_assignment: AttributeAssignment<'db>,
) {
if let Some(class_body_scope) = self.is_method_of_class() {
// We only care about attribute assignments to the first parameter of a method,
// i.e. typically `self` or `cls`.
let accessed_object_refers_to_first_parameter =
object.as_name_expr().map(|name| name.id.as_str())
== self.current_first_parameter_name;
if accessed_object_refers_to_first_parameter {
self.attribute_assignments
.entry(class_body_scope)
.or_default()
.entry(attr.id().clone())
.or_default()
.push(attribute_assignment);
}
}
}
fn add_pattern_constraint(
&mut self,
subject: Expression<'db>,
@ -457,6 +517,20 @@ impl<'db> SemanticIndexBuilder<'db> {
/// Record an expression that needs to be a Salsa ingredient, because we need to infer its type
/// standalone (type narrowing tests, RHS of an assignment.)
fn add_standalone_expression(&mut self, expression_node: &ast::Expr) -> Expression<'db> {
self.add_standalone_expression_impl(expression_node, ExpressionKind::Normal)
}
/// Same as [`SemanticIndexBuilder::add_standalone_expression`], but marks the expression as a
/// *type* expression, which makes sure that it will later be inferred as such.
fn add_standalone_type_expression(&mut self, expression_node: &ast::Expr) -> Expression<'db> {
self.add_standalone_expression_impl(expression_node, ExpressionKind::TypeExpression)
}
fn add_standalone_expression_impl(
&mut self,
expression_node: &ast::Expr,
expression_kind: ExpressionKind,
) -> Expression<'db> {
let expression = Expression::new(
self.db,
self.file,
@ -465,6 +539,7 @@ impl<'db> SemanticIndexBuilder<'db> {
unsafe {
AstNodeRef::new(self.module.clone(), expression_node)
},
expression_kind,
countme::Count::default(),
);
self.expressions_by_node
@ -668,6 +743,11 @@ impl<'db> SemanticIndexBuilder<'db> {
use_def_maps,
imported_modules: Arc::new(self.imported_modules),
has_future_annotations: self.has_future_annotations,
attribute_assignments: self
.attribute_assignments
.into_iter()
.map(|(k, v)| (k, Arc::new(v)))
.collect(),
}
}
}
@ -706,7 +786,17 @@ where
builder.declare_parameters(parameters);
let mut first_parameter_name = parameters
.iter_non_variadic_params()
.next()
.map(|first_param| first_param.parameter.name.id().as_str());
std::mem::swap(
&mut builder.current_first_parameter_name,
&mut first_parameter_name,
);
builder.visit_body(body);
builder.current_first_parameter_name = first_parameter_name;
builder.pop_scope()
},
);
@ -840,6 +930,19 @@ where
unpack: None,
first: false,
}),
ast::Expr::Attribute(ast::ExprAttribute {
value: object,
attr,
..
}) => {
self.register_attribute_assignment(
object,
attr,
AttributeAssignment::Unannotated { value },
);
None
}
_ => None,
};
@ -858,6 +961,7 @@ where
ast::Stmt::AnnAssign(node) => {
debug_assert_eq!(&self.current_assignments, &[]);
self.visit_expr(&node.annotation);
let annotation = self.add_standalone_type_expression(&node.annotation);
if let Some(value) = &node.value {
self.visit_expr(value);
}
@ -869,6 +973,20 @@ where
) {
self.push_assignment(node.into());
self.visit_expr(&node.target);
if let ast::Expr::Attribute(ast::ExprAttribute {
value: object,
attr,
..
}) = &*node.target
{
self.register_attribute_assignment(
object,
attr,
AttributeAssignment::Annotated { annotation },
);
}
self.pop_assignment();
} else {
self.visit_expr(&node.target);