Skip to content

Commit 43b174c

Browse files
authored
[ty] Infer lambda parameter types with Callable type context (astral-sh#24317)
Improves astral-sh#22633 to infer the use of lambda parameters in a lambda body with type context, e.g., ```py x: Callable[[str], str] = lambda x: reveal_type(x) # revealed: str reveal_type(x) # revealed: (x: str) -> str ``` Unlike other definitions, lambda parameter types cannot be determined purely syntactically in semantic indexing. Instead, they depend on the inferred type of the lambda to access its parameter types. Unfortunately, this makes lambda inference cyclic, as the body of the lambda depends on the outer lambda type, and there is no obvious way of splitting out inference of the lambda parameter types from its return type. To avoid initiating cycles on the entire scope containing the lambda, this PR introduces a new inference query — statement-level inference. Statements are a minimal unit of code that encapsulate any internal type context. This makes them very useful to infer a given sub-expression "naturally" without having to provide any external type context. There are other places where we currently rely on scope-level inference for this purpose (e.g., see `infer_complete_scope_types`, the current implementation of astral-sh#23761, and the discussion in astral-sh/ty#3124). Note that statement-level inference is not perfectly fine-grained, e.g., the test expression of an `if` statement does not require external type context and is independent from its body, so statement-level inference may lead to unnecessarily large cycles, but having the unit of code being generalized to an AST structure allows us to avoid the need for such special cases, but this can always change in the future. Additionally, many statements are simply wrappers around definitions or standalone expressions, so we can avoid extra salsa allocations in the common case.
1 parent 4f449ae commit 43b174c

13 files changed

Lines changed: 1041 additions & 133 deletions

File tree

crates/ty_ide/src/semantic_tokens.rs

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ use ruff_python_ast::{
4343
};
4444
use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
4545
use std::ops::Deref;
46-
use ty_python_core::definition::{Definition, DefinitionKind};
46+
use ty_python_core::definition::{Definition, DefinitionKind, ParameterDefinitionNodeKind};
4747
use ty_python_semantic::{
4848
HasType, SemanticModel,
4949
types::ide_support::{
@@ -314,7 +314,7 @@ impl<'db> SemanticTokenVisitor<'db> {
314314
}
315315
DefinitionKind::Class(_) => Some((SemanticTokenType::Class, modifiers)),
316316
DefinitionKind::TypeVar(_) => Some((SemanticTokenType::TypeParameter, modifiers)),
317-
DefinitionKind::Parameter(parameter) => {
317+
DefinitionKind::Parameter(ParameterDefinitionNodeKind::Parameter(parameter)) => {
318318
let parsed = parsed_module(db, definition.file(db));
319319
let ty = parameter.node(&parsed.load(db)).inferred_type(&model);
320320

@@ -342,12 +342,7 @@ impl<'db> SemanticTokenVisitor<'db> {
342342

343343
Some((SemanticTokenType::Parameter, modifiers))
344344
}
345-
DefinitionKind::VariadicPositionalParameter(_) => {
346-
Some((SemanticTokenType::Parameter, modifiers))
347-
}
348-
DefinitionKind::VariadicKeywordParameter(_) => {
349-
Some((SemanticTokenType::Parameter, modifiers))
350-
}
345+
DefinitionKind::Parameter(_) => Some((SemanticTokenType::Parameter, modifiers)),
351346
DefinitionKind::TypeAlias(_) => Some((SemanticTokenType::TypeParameter, modifiers)),
352347
DefinitionKind::Import(_)
353348
| DefinitionKind::ImportFrom(_)

crates/ty_python_core/src/ast_ids.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,12 @@ pub(crate) mod node_key {
150150
}
151151
}
152152

153+
impl From<&ast::ExprLambda> for ExpressionNodeKey {
154+
fn from(value: &ast::ExprLambda) -> Self {
155+
Self(NodeKey::from_node(value))
156+
}
157+
}
158+
153159
impl From<&ast::Identifier> for ExpressionNodeKey {
154160
fn from(value: &ast::Identifier) -> Self {
155161
Self(NodeKey::from_node(value))

0 commit comments

Comments
 (0)