Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[red-knot] Implement support for attributes implicitly declared via their parameter types #16111

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions crates/red_knot_python_semantic/resources/mdtest/attributes.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,20 @@ class C:
if flag:
self.possibly_undeclared_unbound: str = "possibly set in __init__"

param = param if param is None else param + 42
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor:

Maybe something like this for a shorter and more realistic example?

Suggested change
param = param if param is None else param + 42
param = param if param is not None else 0

self.inferred_from_redefined_param = param

def other_method(self, param: str) -> None:
self.inferred_from_param_not_in_init = param

c_instance = C(1)

reveal_type(c_instance.inferred_from_value) # revealed: Unknown | Literal[1, "a"]

# TODO: Same here. This should be `Unknown | Literal[1, "a"]`
reveal_type(c_instance.inferred_from_other_attribute) # revealed: Unknown

# TODO: should be `int | None`
reveal_type(c_instance.inferred_from_param) # revealed: Unknown | int | None
reveal_type(c_instance.inferred_from_param) # revealed: int | None

reveal_type(c_instance.declared_only) # revealed: bytes

Expand All @@ -41,13 +46,17 @@ reveal_type(c_instance.declared_and_bound) # revealed: bool
# mypy and pyright do not show an error here.
reveal_type(c_instance.possibly_undeclared_unbound) # revealed: str

reveal_type(c_instance.inferred_from_redefined_param) # revealed: Unknown | None | int
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for adding this test case. I think we could probably move it directly below the reveal_type(c_instance.inferred_from_param) above (and similarly move the definition in __init__ upwards). And add a small comment that explains why we union with Unknown in that case.


reveal_type(c_instance.inferred_from_param_not_in_init) # revealed: Unknown | str
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test case also makes sense, but we already have that below in the "Variable defined in non-__init__ method" section. Notice that it comes with a TODO that shows that we want to infer the declared parameter type for these cases as well. I think it would be surprising to special-case __init__ in the way you have done here in this PR, although I understand the intention.

So to summarize: I think we can remove this test case here; remove the if name.as_str() == "__init__" check in your implementation, and then resolve the TODO in the test case below.


# This assignment is fine, as we infer `Unknown | Literal[1, "a"]` for `inferred_from_value`.
c_instance.inferred_from_value = "value set on instance"

# This assignment is also fine:
c_instance.inferred_from_param = None

# TODO: this should be an error (incompatible types in assignment)
# error: [invalid-assignment] "Object of type `Literal["incompatible"]` is not assignable to attribute `inferred_from_param` of type `int | None`"
c_instance.inferred_from_param = "incompatible"

# TODO: we already show an error here but the message might be improved?
Expand Down
44 changes: 42 additions & 2 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@ pub(crate) use self::signatures::Signature;
pub use self::subclass_of::SubclassOfType;
use crate::module_name::ModuleName;
use crate::module_resolver::{file_to_module, resolve_module, KnownModule};
use crate::semantic_index::ast_ids::HasScopedExpressionId;
use crate::semantic_index::ast_ids::{HasScopedExpressionId, HasScopedUseId};
use crate::semantic_index::attribute_assignment::AttributeAssignment;
use crate::semantic_index::definition::Definition;
use crate::semantic_index::definition::{Definition, DefinitionKind};
use crate::semantic_index::expression::Expression;
use crate::semantic_index::symbol::{self as symbol, ScopeId, ScopedSymbolId};
use crate::semantic_index::{
Expand Down Expand Up @@ -4240,6 +4240,46 @@ impl<'db> Class<'db> {
//
// self.name = <value>

// Check for a special case - unannotated assignments in `__init__` method
// that assign a method param with declared type. E.g.:
// ```python
// class A:
// def __init__(self, name: str):
// self.name = name
// ```
// In this case we infer attribute type as if it had been declared with
// the type of the value assigned to it, without union with Unknown.
Comment on lines +4243 to +4251
Copy link
Contributor

@sharkdp sharkdp Feb 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe

Suggested change
// Check for a special case - unannotated assignments in `__init__` method
// that assign a method param with declared type. E.g.:
// ```python
// class A:
// def __init__(self, name: str):
// self.name = name
// ```
// In this case we infer attribute type as if it had been declared with
// the type of the value assigned to it, without union with Unknown.
// If we see that a *declared* parameter of a method is directly assigned
// to an instance attribute, we treat that as if the attribute assignment had
// been annotated with that same parameter type. For example, we treat
// the following attribute assignment, as if it had been annotated with
// `self.name: str = name`:
// ```python
// class A:
// def __init__(self, name: str):
// self.name = name
// ```

let value_expr_node = value.node_ref(db).node();

if let ast::Expr::Name(name_expr) = value_expr_node {
let expr_scope_id = value.scope(db);

let use_def_map = use_def_map(db, expr_scope_id);

// Check that the last (and implicitly only) reachable binding of the name
// is in the function definition (parameter declaration).
let bindings =
use_def_map.bindings_at_use(name_expr.scoped_use_id(db, expr_scope_id));

if bindings.last().is_some_and(|binding| {
binding.binding.is_some_and(|definition| {
matches!(definition.kind(db), DefinitionKind::Parameter(_))
&& definition.category(db).is_declaration()
})
}) {
if let Some(ast::StmtFunctionDef { name, .. }) =
expr_scope_id.node(db).as_function()
{
if name.as_str() == "__init__" {
let annotation_ty = infer_expression_type(db, *value);

// TODO: check if there are conflicting declarations
return Symbol::bound(annotation_ty);
Comment on lines +4274 to +4277
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unfortunately, I don't think that this is sufficient.

In the attribute assignment self.name = value, the inferred type for value might be different than the declared type (e.g. due to narrowing). For example, consider:

class C:
    def __init__(self, x: str | None = None) -> None:
        if x is not None:
            self.x = x

reveal_type(C().x)

On this branch, we would get the answer str. This seems fine for the given example, but if we extend that a bit:

class C:
    def __init__(self, x: str | None = None) -> None:
        if x is not None:
            self.x = x
        else:
            self.x = 0

reveal_type(C().x)

we would still get str due to the early return just below, but this is now wrong (too narrow).

I think we should introduce a new section in the attributes.md test suite with a few test cases like the ones above.

To fix this, we might want to retrieve the declared type for the parameter, compare it with the inferred type of the RHS of the assignment, and only trigger this special case if both are equal? This would lead to the answers Unknown | str and Unknown | str | Literal[0] for these examples, which seems good to me.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another test case that we should probably add is the following, to make sure that we infer str and not Literal[""]:

class C:
    def __init__(self, param: str = "") -> None:
        self.x = param

reveal_type(C().x)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the latter case, the inferred and declared type of param will both be str anyway. Defaults do not narrow the inferred type, since we can't assume the default is used for any given call.

(Not that it's a problem to add that test, but I don't think that one would fail even if we looked only at inferred type.)

}
};
Comment on lines +4264 to +4279
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We'll need to wrap this in a salsa query because it access both kind and node, which both are AST nodes from other modules.

}
}

let inferred_ty = infer_expression_type(db, *value);

union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
Expand Down
Loading