diff --git a/crates/red_knot_python_semantic/resources/mdtest/attributes.md b/crates/red_knot_python_semantic/resources/mdtest/attributes.md index 63baa047bc..c51a899131 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/attributes.md +++ b/crates/red_knot_python_semantic/resources/mdtest/attributes.md @@ -232,6 +232,36 @@ reveal_type(c_instance.y) # revealed: int reveal_type(c_instance.z) # revealed: int ``` +#### Attributes defined in multi-target assignments + +```py +class C: + def __init__(self) -> None: + self.a = self.b = 1 + +c_instance = C() + +reveal_type(c_instance.a) # revealed: Unknown | Literal[1] +reveal_type(c_instance.b) # revealed: Unknown | Literal[1] +``` + +#### Augmented assignments + +```py +class Weird: + def __iadd__(self, other: None) -> str: + return "a" + +class C: + def __init__(self) -> None: + self.w = Weird() + self.w += None + +# TODO: Mypy and pyright do not support this, but it would be great if we could +# infer `Unknown | str` or at least `Unknown | Weird | str` here. +reveal_type(C().w) # revealed: Unknown | Weird +``` + #### Attributes defined in tuple unpackings ```py @@ -253,19 +283,24 @@ reveal_type(c_instance.b1) # revealed: Unknown | Literal["a"] reveal_type(c_instance.c1) # revealed: Unknown | int reveal_type(c_instance.d1) # revealed: Unknown | str -# TODO: This should be supported (no error; type should be: `Unknown | Literal[1]`) -# error: [unresolved-attribute] -reveal_type(c_instance.a2) # revealed: Unknown +reveal_type(c_instance.a2) # revealed: Unknown | Literal[1] -# TODO: This should be supported (no error; type should be: `Unknown | Literal["a"]`) -# error: [unresolved-attribute] -reveal_type(c_instance.b2) # revealed: Unknown +reveal_type(c_instance.b2) # revealed: Unknown | Literal["a"] -# TODO: Similar for these two (should be `Unknown | int` and `Unknown | str`, respectively) -# error: [unresolved-attribute] -reveal_type(c_instance.c2) # revealed: Unknown -# error: [unresolved-attribute] -reveal_type(c_instance.d2) # revealed: Unknown +reveal_type(c_instance.c2) # revealed: Unknown | int +reveal_type(c_instance.d2) # revealed: Unknown | str +``` + +#### Starred assignments + +```py +class C: + def __init__(self) -> None: + self.a, *self.b = (1, 2, 3) + +c_instance = C() +reveal_type(c_instance.a) # revealed: Unknown | Literal[1] +reveal_type(c_instance.b) # revealed: Unknown | @Todo(starred unpacking) ``` #### Attributes defined in for-loop (unpacking) @@ -287,6 +322,8 @@ class TupleIterable: def __iter__(self) -> TupleIterator: return TupleIterator() +class NonIterable: ... + class C: def __init__(self): for self.x in IntIterable(): @@ -295,14 +332,54 @@ class C: for _, self.y in TupleIterable(): pass -# TODO: Pyright fully supports these, mypy detects the presence of the attributes, -# but infers type `Any` for both of them. We should infer `int` and `str` here: + # TODO: We should emit a diagnostic here + for self.z in NonIterable(): + pass -# error: [unresolved-attribute] -reveal_type(C().x) # revealed: Unknown +reveal_type(C().x) # revealed: Unknown | int +reveal_type(C().y) # revealed: Unknown | str +``` + +#### Attributes defined in `with` statements + +```py +class ContextManager: + def __enter__(self) -> int | None: ... + def __exit__(self, exc_type, exc_value, traceback) -> None: ... + +class C: + def __init__(self) -> None: + with ContextManager() as self.x: + pass + +c_instance = C() + +# TODO: Should be `Unknown | int | None` # error: [unresolved-attribute] -reveal_type(C().y) # revealed: Unknown +reveal_type(c_instance.x) # revealed: Unknown +``` + +#### Attributes defined in comprehensions + +```py +class IntIterator: + def __next__(self) -> int: + return 1 + +class IntIterable: + def __iter__(self) -> IntIterator: + return IntIterator() + +class C: + def __init__(self) -> None: + [... for self.a in IntIterable()] + +c_instance = C() + +# TODO: Should be `Unknown | int` +# error: [unresolved-attribute] +reveal_type(c_instance.a) # revealed: Unknown ``` #### Conditionally declared / bound attributes @@ -443,6 +520,15 @@ class C: reveal_type(C().x) # revealed: str ``` +#### Diagnostics are reported for the right-hand side of attribute assignments + +```py +class C: + def __init__(self) -> None: + # error: [too-many-positional-arguments] + self.x: int = len(1, 2, 3) +``` + ### Pure class variables (`ClassVar`) #### Annotated with `ClassVar` type qualifier diff --git a/crates/red_knot_python_semantic/src/semantic_index/attribute_assignment.rs b/crates/red_knot_python_semantic/src/semantic_index/attribute_assignment.rs index 29f9a36d12..53cc41e9dc 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/attribute_assignment.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/attribute_assignment.rs @@ -1,4 +1,7 @@ -use crate::semantic_index::expression::Expression; +use crate::{ + semantic_index::{ast_ids::ScopedExpressionId, expression::Expression}, + unpack::Unpack, +}; use ruff_python_ast::name::Name; @@ -14,6 +17,17 @@ pub(crate) enum AttributeAssignment<'db> { /// An attribute assignment without a type annotation, e.g. `self.x = `. Unannotated { value: Expression<'db> }, + + /// An attribute assignment where the right-hand side is an iterable, for example + /// `for self.x in `. + Iterable { iterable: Expression<'db> }, + + /// An attribute assignment where the left-hand side is an unpacking expression, + /// e.g. `self.x, self.y = `. + Unpack { + attribute_expression_id: ScopedExpressionId, + unpack: Unpack<'db>, + }, } pub(crate) type AttributeAssignments<'db> = FxHashMap>>; diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index ad0655818a..fff95086c1 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -6,9 +6,9 @@ use rustc_hash::{FxHashMap, FxHashSet}; use ruff_db::files::File; use ruff_db::parsed::ParsedModule; use ruff_index::IndexVec; -use ruff_python_ast as ast; use ruff_python_ast::name::Name; use ruff_python_ast::visitor::{walk_expr, walk_pattern, walk_stmt, Visitor}; +use ruff_python_ast::{self as ast, ExprContext}; use crate::ast_node_ref::AstNodeRef; use crate::module_name::ModuleName; @@ -1231,6 +1231,20 @@ where unpack: None, first: false, }), + ast::Expr::Attribute(ast::ExprAttribute { + value: object, + attr, + .. + }) => { + self.register_attribute_assignment( + object, + attr, + AttributeAssignment::Iterable { + iterable: iter_expr, + }, + ); + None + } _ => None, }; @@ -1459,7 +1473,7 @@ where fn visit_expr(&mut self, expr: &'ast ast::Expr) { self.scopes_by_expression .insert(expr.into(), self.current_scope()); - self.current_ast_ids().record_expression(expr); + let expression_id = self.current_ast_ids().record_expression(expr); match expr { ast::Expr::Name(name_node @ ast::ExprName { id, ctx, .. }) => { @@ -1718,6 +1732,35 @@ where self.simplify_visibility_constraints(pre_op); } + ast::Expr::Attribute(ast::ExprAttribute { + value: object, + attr, + ctx: ExprContext::Store, + range: _, + }) => { + if let Some( + CurrentAssignment::Assign { + unpack: Some(unpack), + .. + } + | CurrentAssignment::For { + unpack: Some(unpack), + .. + }, + ) = self.current_assignment() + { + self.register_attribute_assignment( + object, + attr, + AttributeAssignment::Unpack { + attribute_expression_id: expression_id, + unpack, + }, + ); + } + + walk_expr(self, expr); + } _ => { walk_expr(self, expr); } diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index 981d6f20ad..360b9a5546 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/crates/red_knot_python_semantic/src/types.rs @@ -40,6 +40,7 @@ use crate::types::call::{ }; use crate::types::class_base::ClassBase; use crate::types::diagnostic::INVALID_TYPE_FORM; +use crate::types::infer::infer_unpack_types; use crate::types::mro::{Mro, MroError, MroIterator}; use crate::types::narrow::narrowing_constraint; use crate::{Db, FxOrderSet, Module, Program, PythonVersion}; @@ -4231,6 +4232,33 @@ impl<'db> Class<'db> { union_of_inferred_types = union_of_inferred_types.add(inferred_ty); } + AttributeAssignment::Iterable { iterable } => { + // We found an attribute assignment like: + // + // for self.name in : + + // TODO: Potential diagnostics resulting from the iterable are currently not reported. + + let iterable_ty = infer_expression_type(db, *iterable); + let inferred_ty = iterable_ty.iterate(db).unwrap_without_diagnostic(); + + union_of_inferred_types = union_of_inferred_types.add(inferred_ty); + } + AttributeAssignment::Unpack { + attribute_expression_id, + unpack, + } => { + // We found an unpacking assignment like: + // + // .., self.name, .. = + // (.., self.name, ..) = + // [.., self.name, ..] = + + let inferred_ty = infer_unpack_types(db, *unpack) + .get(*attribute_expression_id) + .expect("Failed to look up type of attribute in unpack assignment"); + union_of_inferred_types = union_of_inferred_types.add(inferred_ty); + } } } diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 7de554f538..6f224b398d 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -200,7 +200,7 @@ pub(crate) fn infer_expression_types<'db>( /// type of the variables involved in this unpacking along with any violations that are detected /// during this unpacking. #[salsa::tracked(return_ref)] -fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> UnpackResult<'db> { +pub(super) fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> UnpackResult<'db> { let file = unpack.file(db); let _span = tracing::trace_span!("infer_unpack_types", range=?unpack.range(db), file=%file.path(db)) diff --git a/crates/red_knot_python_semantic/src/types/unpacker.rs b/crates/red_knot_python_semantic/src/types/unpacker.rs index d1b288ed18..f7782d4e6b 100644 --- a/crates/red_knot_python_semantic/src/types/unpacker.rs +++ b/crates/red_knot_python_semantic/src/types/unpacker.rs @@ -72,11 +72,9 @@ impl<'db> Unpacker<'db> { value_ty: Type<'db>, ) { match target { - ast::Expr::Name(target_name) => { - self.targets.insert( - target_name.scoped_expression_id(self.db(), self.scope), - value_ty, - ); + ast::Expr::Name(_) | ast::Expr::Attribute(_) => { + self.targets + .insert(target.scoped_expression_id(self.db(), self.scope), value_ty); } ast::Expr::Starred(ast::ExprStarred { value, .. }) => { self.unpack_inner(value, value_expr, value_ty);