Skip to content

Comments

[red-knot] infer attribute assignments bound in comprehensions#17396

Merged
dhruvmanila merged 8 commits intoastral-sh:mainfrom
mtshiba:comprehension-attribute-assignment
Apr 19, 2025
Merged

[red-knot] infer attribute assignments bound in comprehensions#17396
dhruvmanila merged 8 commits intoastral-sh:mainfrom
mtshiba:comprehension-attribute-assignment

Conversation

@mtshiba
Copy link
Collaborator

@mtshiba mtshiba commented Apr 14, 2025

Summary

This PR is a follow-up to #16852.

Instance variables bound in comprehensions are recorded, allowing type inference to work correctly.

This required adding support for unpacking in comprehension which resolves #15369.

Test Plan

One TODO in mdtest/attributes.md is now resolved, and some new test cases are added.

@github-actions
Copy link
Contributor

github-actions bot commented Apr 14, 2025

mypy_primer results

Changes were detected when running on open source projects
async-utils (https://github.com/mikeshardmind/async-utils)
- error[lint:unsupported-operator] /tmp/mypy_primer/projects/async-utils/_misc/_queue_testing.py:50:30: Operator `+` is unsupported between objects of type `Literal["\u{1b}[30;1m%(asctime)s\u{1b}[0m "]` and `tuple[Unknown, Literal["\u{1b}[40;1m"]] | tuple[Unknown, Literal["\u{1b}[34;1m"]] | tuple[Unknown, Literal["\u{1b}[33;1m"]] | tuple[Unknown, Literal["\u{1b}[31m"]] | tuple[Unknown, Literal["\u{1b}[41m"]]`
- Found 30 diagnostics
+ Found 29 diagnostics

@AlexWaygood AlexWaygood added the ty Multi-file analysis & type inference label Apr 14, 2025
@dhruvmanila
Copy link
Member

By the looks of it, this PR resolves #15369 ?

Copy link
Contributor

@carljm carljm left a comment

Choose a reason for hiding this comment

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

This looks really good!

As @dhruvmanila mentions, it seems like this adds general support for unpacking in comprehension targets, not only for discovering instance attributes? If so, can we add some tests specifically for unpacking in comprehensions, to regular Name targets?

@AlexWaygood AlexWaygood removed their request for review April 15, 2025 16:17
@mtshiba mtshiba requested review from carljm and dhruvmanila April 15, 2025 16:18
Copy link
Member

@dhruvmanila dhruvmanila left a comment

Choose a reason for hiding this comment

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

This looks great. My comments are mainly around adding / updating comments and increasing the test coverage for unpacking in comprehensions.

Comment on lines 754 to 760
target: AstNodeRef<ast::ExprName>,
target: AstNodeRef<ast::Expr>,
Copy link
Member

Choose a reason for hiding this comment

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

Unrelated to this PR but I think the target here is either a ExprName or ExprAttribute, so it might be useful to have an enum similar to TargetKind which has two variants like TargetExpr::Name(ast::ExprName) and TargetExpr::Attribute(ast::ExprAttribute). This should not be done in this PR as it requires updating other attribute unpacking as well.

Copy link
Collaborator Author

@mtshiba mtshiba Apr 16, 2025

Choose a reason for hiding this comment

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

OK, I'll work on this in a separate PR (this would be a topic related to #16894).

We also need to support ExprSubscript as a target expression.

Copy link
Member

Choose a reason for hiding this comment

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

Just as a note it's not a priority to make this refactor, just something nice to have and is not necessarily required.

@sharkdp sharkdp removed their request for review April 15, 2025 18:19
@mtshiba mtshiba force-pushed the comprehension-attribute-assignment branch from 7ec0755 to 01c959e Compare April 16, 2025 16:12
@mtshiba mtshiba force-pushed the comprehension-attribute-assignment branch from 01c959e to 86fc13e Compare April 16, 2025 16:22
@mtshiba mtshiba requested a review from MichaReiser as a code owner April 16, 2025 16:22
@mtshiba mtshiba requested a review from dhruvmanila April 16, 2025 16:29
Copy link
Contributor

@carljm carljm left a comment

Choose a reason for hiding this comment

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

This looks good to me, but will let @dhruvmanila have another look and merge if happy with it.

Comment on lines 3903 to 3919
// TODO: support attribute, subscription
fn infer_comprehension_target(&mut self, target: &ast::Expr) {
match target {
ast::Expr::Name(name) => {
self.infer_definition(name);
}
ast::Expr::Tuple(ast::ExprTuple { elts, .. })
| ast::Expr::List(ast::ExprList { elts, .. }) => {
for elt in elts {
self.infer_comprehension_target(elt);
}
}
_ => {
self.infer_expression(target);
}
}
}
Copy link
Member

Choose a reason for hiding this comment

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

Shouldn't this just use the existing infer_target method? Or, why is a separate method required in this case?

Copy link
Member

@dhruvmanila dhruvmanila Apr 17, 2025

Choose a reason for hiding this comment

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

That method (infer_target) isn't correct currently as it has a bug (https://github.com/astral-sh/ruff/issues/16514) but we're aware of it.

Copy link
Collaborator Author

@mtshiba mtshiba Apr 18, 2025

Choose a reason for hiding this comment

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

I think we cannot simply replace infer_comprehension_target with infer_target.

Specifically, if we change the code as follows, the tests will fail as follows:

# infer.rs#L3897
-         self.infer_comprehension_target(target);
+         self.infer_target(target, iter, |db, iter_ty| {
+             iter_ty.iterate(db)
+         });
invalid.md - Tests for invalid types in type expressions - Invalid Collection based AST nodes

  crates\red_knot_python_semantic\resources\mdtest\annotations\invalid.md:98 panicked at crates\red_knot_python_semantic\src\types\infer.rs:561:9
  crates\red_knot_python_semantic\resources\mdtest\annotations\invalid.md:98 assertion `left == right` failed
  left: ScopeId { [salsa id]: Id(801), file: File(System("\\src\\mdtest_snippet.py")), file_scope_id: FileScopeId(1), count: Count { ghost: PhantomData<fn(red_knot_python_semantic::semantic_index::symbol::ScopeId)> } }
 right: ScopeId { [salsa id]: Id(800), file: File(System("\\src\\mdtest_snippet.py")), file_scope_id: FileScopeId(0), count: Count { ghost: PhantomData<fn(red_knot_python_semantic::semantic_index::symbol::ScopeId)> } }

---

attributes.md - Attributes - Class and instance variables - Pure instance variables - Attributes defined in comprehensions

  crates\red_knot_python_semantic\resources\mdtest\attributes.md:392 panicked at crates\red_knot_python_semantic\src\types\infer.rs:561:9
  crates\red_knot_python_semantic\resources\mdtest\attributes.md:392 assertion `left == right` failed
  left: ScopeId { [salsa id]: Id(37ac), file: File(System("\\src\\mdtest_snippet.py")), file_scope_id: FileScopeId(11), count: Count { ghost: PhantomData<fn(red_knot_python_semantic::semantic_index::symbol::ScopeId)> } }
 right: ScopeId { [salsa id]: Id(37a9), file: File(System("\\src\\mdtest_snippet.py")), file_scope_id: FileScopeId(10), count: Count { ghost: PhantomData<fn(red_knot_python_semantic::semantic_index::symbol::ScopeId)> } }

TypeInferenceBuilder::extend fails because the scope of the comprehension and the scope of the value=iter are different.

However, since the implementations of infer_comprehension_target and infer_target_impl are almost the same, I will use this instead.

It would be better if infer_target supported comprehension.

Copy link
Member

Choose a reason for hiding this comment

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

It would be better if infer_target supported comprehension.

Agreed. Thanks for making the change. I think the logic to infer the value expression using infer_same_file_expression_type is only to be done for the first comprehension, right? I think we could update the infer_target specifically the to_assigned_ty to instead be infer_value_expr so that the closure is responsible to infer the (non Name node) value expression so that we can check whether is_first is true inside the infer_comprehension.

I do have to change the logic of infer_target a bit to fix the bug as mentioned in https://github.com/astral-sh/ruff/issues/16514 but that's something for later.

Copy link
Member

Choose a reason for hiding this comment

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

I've pushed 329a0f2 (#17396) to reflect what I propose.

dhruvmanila added a commit that referenced this pull request Apr 18, 2025
## Summary

Part of #15383, this PR adds support for overloaded callables.

Typing spec: https://typing.python.org/en/latest/spec/overload.html

Specifically, it does the following:
1. Update the `FunctionType::signature` method to return signatures from
a possibly overloaded callable using a new `FunctionSignature` enum
2. Update `CallableType` to accommodate overloaded callable by updating
the inner type to `Box<[Signature]>`
3. Update the relation methods on `CallableType` with logic specific to
overloads
4. Update the display of callable type to display a list of signatures
enclosed by parenthesis
5. Update `CallableTypeOf` special form to recognize overloaded callable
6. Update subtyping, assignability and fully static check to account for
callables (equivalence is planned to be done as a follow-up)

For (2), it is required to be done in this PR because otherwise I'd need
to add some workaround for `into_callable_type` and I though it would be
best to include it in here.

For (2), another possible design would be convert `CallableType` in an
enum with two variants `CallableType::Single` and
`CallableType::Overload` but I decided to go with `Box<[Signature]>` for
now to (a) mirror it to be equivalent to `overload` field on
`CallableSignature` and (b) to avoid any refactor in this PR. This could
be done in a follow-up to better split the two kind of callables.

### Design

There were two main candidates on how to represent the overloaded
definition:
1. To include it in the existing infrastructure which is what this PR is
doing by recognizing all the signatures within the
`FunctionType::signature` method
2. To create a new `Overload` type variant

<details><summary>For context, this is what I had in mind with the new
type variant:</summary>
<p>

```rs
pub enum Type {
	FunctionLiteral(FunctionType),
    Overload(OverloadType),
    BoundMethod(BoundMethodType),
    ...
}

pub struct OverloadType {
	// FunctionLiteral or BoundMethod
    overloads: Box<[Type]>,
	// FunctionLiteral or BoundMethod
    implementation: Option<Type>
}

pub struct BoundMethodType {
    kind: BoundMethodKind,
    self_instance: Type,
}

pub enum BoundMethodKind {
    Function(FunctionType),
    Overload(OverloadType),
}
```

</p>
</details> 

The main reasons to choose (1) are the simplicity in the implementation,
reusing the existing infrastructure, avoiding any complications that the
new type variant has specifically around the different variants between
function and methods which would require the overload type to use `Type`
instead.

### Implementation

The core logic is how to collect all the overloaded functions. The way
this is done in this PR is by recording a **use** on the `Identifier`
node that represents the function name in the use-def map. This is then
used to fetch the previous symbol using the same name. This way the
signatures are going to be propagated from top to bottom (from first
overload to the final overload or the implementation) with each function
/ method. For example:

```py
from typing import overload

@overload
def foo(x: int) -> int: ...
@overload
def foo(x: str) -> str: ...
def foo(x: int | str) -> int | str:
	return x
```

Here, each definition of `foo` knows about all the signatures that comes
before itself. So, the first overload would only see itself, the second
would see the first and itself and so on until the implementation or the
final overload.

This approach required some updates specifically recognizing
`Identifier` node to record the function use because it doesn't use
`ExprName`.

## Test Plan

Update existing test cases which were limited by the overload support
and add test cases for the following cases:
* Valid overloads as functions, methods, generics, version specific
* Invalid overloads as stated in
https://typing.python.org/en/latest/spec/overload.html#invalid-overload-definitions
(implementation will be done in a follow-up)
* Various relation: fully static, subtyping, and assignability (others
in a follow-up)

## Ecosystem changes

_WIP_

After going through the ecosystem changes (there are a lot!), here's
what I've found:

We need assignability check between a callable type and a class literal
because a lot of builtins are defined as classes in typeshed whose
constructor method is overloaded e.g., `map`, `sorted`, `list.sort`,
`max`, `min` with the `key` parameter, `collections.abc.defaultdict`,
etc. (https://github.com/astral-sh/ruff/issues/17343). This makes up
most of the ecosystem diff **roughly 70 diagnostics**. For example:

```py
from collections import defaultdict

# red-knot: No overload of bound method `__init__` matches arguments [lint:no-matching-overload]
defaultdict(int)
# red-knot: No overload of bound method `__init__` matches arguments [lint:no-matching-overload]
defaultdict(list)

class Foo:
    def __init__(self, x: int):
        self.x = x

# red-knot: No overload of function `__new__` matches arguments [lint:no-matching-overload]
map(Foo, ["a", "b", "c"])
```

Duplicate diagnostics in unpacking
(https://github.com/astral-sh/ruff/issues/16514) has **~16
diagnostics**.

Support for the `callable` builtin which requires `TypeIs` support. This
is **5 diagnostics**. For example:
```py
from typing import Any

def _(x: Any | None) -> None:
    if callable(x):
        # red-knot: `Any | None`
        # Pyright: `(...) -> object`
        # mypy: `Any`
        # pyrefly: `(...) -> object`
        reveal_type(x)
```

Narrowing on `assert` which has **11 diagnostics**. This is being worked
on in #17345. For example:
```py
import re

match = re.search("", "")
assert match
match.group()  # error: [possibly-unbound-attribute]
```

Others:
* `Self`: 2
* Type aliases: 6
* Generics: 3
* Protocols: 13
* Unpacking in comprehension: 1
(#17396)

## Performance

Refer to
#17366 (comment).
@mtshiba mtshiba force-pushed the comprehension-attribute-assignment branch from 8616291 to 67d1150 Compare April 18, 2025 17:37
@dhruvmanila
Copy link
Member

The mypy_primer change is removing a false positive, it popped up in the overload PR.

Copy link
Member

@dhruvmanila dhruvmanila left a comment

Choose a reason for hiding this comment

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

Thank you!

@dhruvmanila dhruvmanila merged commit da6b68c into astral-sh:main Apr 19, 2025
22 checks passed
dcreager added a commit that referenced this pull request Apr 21, 2025
* main:
  Update pre-commit dependencies (#17506)
  [red-knot] Simplify visibility constraint handling for `*`-import definitions (#17486)
  [red-knot] Detect (some) invalid protocols (#17488)
  [red-knot] Correctly identify protocol classes (#17487)
  Update dependency ruff to v0.11.6 (#17516)
  Update Rust crate shellexpand to v3.1.1 (#17512)
  Update Rust crate proc-macro2 to v1.0.95 (#17510)
  Update Rust crate rand to v0.9.1 (#17511)
  Update Rust crate libc to v0.2.172 (#17509)
  Update Rust crate jiff to v0.2.9 (#17508)
  Update Rust crate clap to v4.5.37 (#17507)
  Update astral-sh/setup-uv action to v5.4.2 (#17504)
  Update taiki-e/install-action digest to 09dc018 (#17503)
  [red-knot] infer attribute assignments bound in comprehensions (#17396)
  [red-knot] simplify gradually-equivalent types out of unions and intersections (#17467)
  [red-knot] pull primer projects to run from file (#17473)
@mtshiba mtshiba deleted the comprehension-attribute-assignment branch April 27, 2025 09:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ty Multi-file analysis & type inference

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[red-knot] Add support for unpacking targets in comprehensions

4 participants