Skip to content

Commit 7eda4ee

Browse files
ibraheemdevibraheemdev
committed
infer lambda parameter types with Callable type context
1 parent 3a95cd9 commit 7eda4ee

7 files changed

Lines changed: 98 additions & 45 deletions

File tree

‎crates/ty_ide/src/semantic_tokens.rs‎

Lines changed: 7 additions & 12 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_semantic::semantic_index::definition::Definition;
46+
use ty_python_semantic::semantic_index::definition::{Definition, ParameterDefinitionNodeKind};
4747
use ty_python_semantic::types::TypeVarKind;
4848
use ty_python_semantic::{
4949
HasType, SemanticModel, semantic_index::definition::DefinitionKind, types::Type,
@@ -308,7 +308,7 @@ impl<'db> SemanticTokenVisitor<'db> {
308308
}
309309
DefinitionKind::Class(_) => Some((SemanticTokenType::Class, modifiers)),
310310
DefinitionKind::TypeVar(_) => Some((SemanticTokenType::TypeParameter, modifiers)),
311-
DefinitionKind::Parameter(parameter) => {
311+
DefinitionKind::Parameter(ParameterDefinitionNodeKind::Parameter(parameter)) => {
312312
let parsed = parsed_module(db, definition.file(db));
313313
let ty = parameter.node(&parsed.load(db)).inferred_type(&model);
314314

@@ -336,12 +336,7 @@ impl<'db> SemanticTokenVisitor<'db> {
336336

337337
Some((SemanticTokenType::Parameter, modifiers))
338338
}
339-
DefinitionKind::VariadicPositionalParameter(_) => {
340-
Some((SemanticTokenType::Parameter, modifiers))
341-
}
342-
DefinitionKind::VariadicKeywordParameter(_) => {
343-
Some((SemanticTokenType::Parameter, modifiers))
344-
}
339+
DefinitionKind::Parameter(_) => Some((SemanticTokenType::Parameter, modifiers)),
345340
DefinitionKind::TypeAlias(_) => Some((SemanticTokenType::TypeParameter, modifiers)),
346341
DefinitionKind::Import(_)
347342
| DefinitionKind::ImportFrom(_)
@@ -1787,7 +1782,7 @@ b: list["int | str"] | None
17871782
c: "list[int | str] | None"
17881783
d: "list[int | str]" | "None"
17891784
e: 'list["int | str"] | "None"'
1790-
f: """'list["int | str"]' | 'None'"""
1785+
f: """'list["int | str"]' | 'None'"""
17911786
"#,
17921787
);
17931788

@@ -2148,7 +2143,7 @@ class Baz:
21482143
prop: str = \"hello\"
21492144
21502145
baz = Baz()
2151-
s = baz.method
2146+
s = baz.method
21522147
t = baz.CONSTANT
21532148
r = baz.prop
21542149
q = Baz.prop
@@ -2383,7 +2378,7 @@ class MyClass:
23832378
def __init__(self): pass
23842379
23852380
"""unrelated string"""
2386-
2381+
23872382
x: str = "hello"
23882383
"#,
23892384
);
@@ -2414,7 +2409,7 @@ What a good module wooo
24142409
def my_func(): pass
24152410
24162411
"""unrelated string"""
2417-
2412+
24182413
x: str = "hello"
24192414
"#,
24202415
);

‎crates/ty_python_semantic/resources/mdtest/bidirectional.md‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -419,9 +419,8 @@ reveal_type(f2) # revealed: (x: int) -> Bar
419419
f3: Callable[[int], Bar] = lambda x: {}
420420
reveal_type(f3) # revealed: (int, /) -> Bar
421421

422-
# TODO: This should reveal `str`.
423-
f4: Callable[[str], str] = lambda x: reveal_type(x) # revealed: Unknown
424-
reveal_type(f4) # revealed: (x: str) -> Unknown
422+
f4: Callable[[str], str] = lambda x: reveal_type(x) # revealed: str
423+
reveal_type(f4) # revealed: (x: str) -> str
425424

426425
# TODO: This should not error once we support `Unpack`.
427426
# error: [invalid-assignment]
@@ -439,15 +438,27 @@ reveal_type(f7) # revealed: (int, /) -> None
439438
# TODO: This should reveal `(*args: int, *, x=1) -> None` once we support `Unpack`.
440439
f8: Callable[[*tuple[int, ...], int], None] = lambda *args, x=1: None
441440
reveal_type(f8) # revealed: (*args, *, x=1) -> None
441+
442+
f9: Callable[[str, int, str], tuple[str, int, str]] = lambda x, y, z: reveal_type((x, y, z)) # revealed: tuple[str, int, str]
443+
reveal_type(f9) # revealed: (x: str, y: int, z: str) -> tuple[str, int, str]
444+
445+
# TODO: This should reveal `tuple[int, ...]` once we support `Unpack`.
446+
f10: Callable[[*tuple[int, ...]], tuple[int, ...]] = lambda *args: reveal_type(args) # revealed: tuple[Unknown, ...]
447+
reveal_type(f10) # revealed: (*args) -> tuple[Unknown, ...]
448+
449+
# TODO: Better generic call inference.
450+
def _(x: list[int]):
451+
f11 = list(map(lambda y: y + 1, x))
452+
reveal_type(f11) # revealed: list[Unknown]
442453
```
443454

444455
We do not currently account for type annotations present later in the scope:
445456

446457
```py
447-
f9 = lambda: [1]
458+
f12 = lambda: [1]
448459
# TODO: This should not error.
449-
_: list[int | str] = f9() # error: [invalid-assignment]
450-
reveal_type(f9) # revealed: () -> list[int]
460+
_: list[int | str] = f12() # error: [invalid-assignment]
461+
reveal_type(f12) # revealed: () -> list[int]
451462
```
452463

453464
## Dunder Calls

‎crates/ty_python_semantic/src/semantic_index/ast_ids.rs‎

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

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

‎crates/ty_python_semantic/src/semantic_index/builder.rs‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1673,14 +1673,18 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
16731673
debug_assert_eq!(existing_definition, None);
16741674
}
16751675

1676-
fn declare_lambda_parameters(&mut self, parameters: &'ast ast::Parameters) {
1676+
fn declare_lambda_parameters(
1677+
&mut self,
1678+
parameters: &'ast ast::Parameters,
1679+
lambda: &'ast ast::ExprLambda,
1680+
) {
16771681
let mut index = 0;
16781682
for parameter in &parameters.posonlyargs {
1679-
self.declare_lambda_parameter(index, parameter);
1683+
self.declare_lambda_parameter(index, parameter, lambda);
16801684
index += 1;
16811685
}
16821686
for parameter in &parameters.args {
1683-
self.declare_lambda_parameter(index, parameter);
1687+
self.declare_lambda_parameter(index, parameter, lambda);
16841688
index += 1;
16851689
}
16861690
if let Some(vararg) = parameters.vararg.as_ref() {
@@ -1692,13 +1696,14 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
16921696
symbol.into(),
16931697
LambdaParameterDefinitionNodeRef {
16941698
index,
1699+
lambda,
16951700
parameter: ParameterDefinitionNodeRef::VariadicPositionalParameter(vararg),
16961701
},
16971702
);
16981703
index += 1;
16991704
}
17001705
for parameter in &parameters.kwonlyargs {
1701-
self.declare_lambda_parameter(index, parameter);
1706+
self.declare_lambda_parameter(index, parameter, lambda);
17021707
index += 1;
17031708
}
17041709
if let Some(kwarg) = parameters.kwarg.as_ref() {
@@ -1710,6 +1715,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
17101715
symbol.into(),
17111716
LambdaParameterDefinitionNodeRef {
17121717
index,
1718+
lambda,
17131719
parameter: ParameterDefinitionNodeRef::VariadicKeywordParameter(kwarg),
17141720
},
17151721
);
@@ -1720,13 +1726,15 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
17201726
&mut self,
17211727
index: usize,
17221728
parameter: &'ast ast::ParameterWithDefault,
1729+
lambda: &'ast ast::ExprLambda,
17231730
) {
17241731
let symbol = self.add_symbol(parameter.name().id().clone());
17251732

17261733
let definition = self.add_definition(
17271734
symbol.into(),
17281735
LambdaParameterDefinitionNodeRef {
17291736
index,
1737+
lambda,
17301738
parameter: ParameterDefinitionNodeRef::Parameter(parameter),
17311739
},
17321740
);
@@ -3244,7 +3252,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
32443252

32453253
// Add symbols and definitions for the parameters to the lambda scope.
32463254
if let Some(parameters) = lambda.parameters.as_ref() {
3247-
self.declare_lambda_parameters(parameters);
3255+
self.declare_lambda_parameters(parameters, lambda);
32483256
}
32493257

32503258
self.visit_expr(lambda.body.as_ref());

‎crates/ty_python_semantic/src/semantic_index/definition.rs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -392,13 +392,13 @@ impl<'ast, 'db> From<ComprehensionDefinitionNodeRef<'ast, 'db>> for DefinitionNo
392392
}
393393
}
394394

395-
impl<'ast, 'db> From<ParameterDefinitionNodeRef<'ast>> for DefinitionNodeRef<'ast, 'db> {
395+
impl<'ast> From<ParameterDefinitionNodeRef<'ast>> for DefinitionNodeRef<'ast, '_> {
396396
fn from(node: ParameterDefinitionNodeRef<'ast>) -> Self {
397397
Self::Parameter(node)
398398
}
399399
}
400400

401-
impl<'ast, 'db> From<LambdaParameterDefinitionNodeRef<'ast>> for DefinitionNodeRef<'ast, 'db> {
401+
impl<'ast> From<LambdaParameterDefinitionNodeRef<'ast>> for DefinitionNodeRef<'ast, '_> {
402402
fn from(node: LambdaParameterDefinitionNodeRef<'ast>) -> Self {
403403
Self::LambdaParameter(node)
404404
}
@@ -546,8 +546,9 @@ impl ParameterDefinitionNodeRef<'_> {
546546

547547
#[derive(Copy, Clone, Debug)]
548548
pub(crate) struct LambdaParameterDefinitionNodeRef<'ast> {
549-
index: usize,
550-
parameter: ParameterDefinitionNodeRef<'ast>,
549+
pub(crate) index: usize,
550+
pub(crate) parameter: ParameterDefinitionNodeRef<'ast>,
551+
pub(crate) lambda: &'ast ast::ExprLambda,
551552
}
552553

553554
#[derive(Copy, Clone, Debug)]
@@ -671,9 +672,11 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
671672
DefinitionNodeRef::LambdaParameter(LambdaParameterDefinitionNodeRef {
672673
index,
673674
parameter,
675+
lambda,
674676
}) => DefinitionKind::LambdaParameter(LambdaParameterDefinitionNodeKind {
675677
index,
676678
parameter: parameter.into_owned(parsed),
679+
lambda: AstNodeRef::new(parsed, lambda),
677680
}),
678681
DefinitionNodeRef::WithItem(WithItemDefinitionNodeRef {
679682
unpack,
@@ -1228,6 +1231,7 @@ impl ParameterDefinitionNodeKind {
12281231
#[derive(Clone, Debug, get_size2::GetSize)]
12291232
pub struct LambdaParameterDefinitionNodeKind {
12301233
pub(crate) index: usize,
1234+
pub(crate) lambda: AstNodeRef<ast::ExprLambda>,
12311235
pub(crate) parameter: ParameterDefinitionNodeKind,
12321236
}
12331237

‎crates/ty_python_semantic/src/types/infer/builder.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -824,31 +824,34 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
824824
}
825825
DefinitionKind::LambdaParameter(LambdaParameterDefinitionNodeKind {
826826
index,
827+
lambda,
827828
parameter: ParameterDefinitionNodeKind::VariadicPositionalParameter(parameter),
828829
}) => {
829830
self.infer_variadic_positional_lambda_parameter_definition(
830831
*index,
831832
parameter.node(self.module()),
833+
lambda.node(self.module()),
832834
definition,
833835
);
834836
}
835837
DefinitionKind::LambdaParameter(LambdaParameterDefinitionNodeKind {
836-
index,
837838
parameter: ParameterDefinitionNodeKind::VariadicKeywordParameter(parameter),
839+
..
838840
}) => {
839841
self.infer_variadic_keyword_lambda_parameter_definition(
840-
*index,
841842
parameter.node(self.module()),
842843
definition,
843844
);
844845
}
845846
DefinitionKind::LambdaParameter(LambdaParameterDefinitionNodeKind {
846847
index,
848+
lambda,
847849
parameter: ParameterDefinitionNodeKind::Parameter(parameter_with_default),
848850
}) => {
849851
self.infer_lambda_parameter_definition(
850852
*index,
851853
parameter_with_default.node(self.module()),
854+
lambda.node(self.module()),
852855
definition,
853856
);
854857
}

‎crates/ty_python_semantic/src/types/infer/builder/function.rs‎

Lines changed: 42 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ use crate::{
2828
},
2929
function_known_decorators, nearest_enclosing_function,
3030
},
31-
infer_definition_types, infer_scope_types, todo_type,
31+
infer_complete_scope_types, infer_definition_types, infer_scope_types, todo_type,
3232
},
3333
};
3434

@@ -939,8 +939,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
939939
/// in the lambda body scope.
940940
pub(super) fn infer_lambda_parameter_definition(
941941
&mut self,
942-
_index: usize,
942+
index: usize,
943943
parameter_with_default: &'ast ast::ParameterWithDefault,
944+
lambda: &'ast ast::ExprLambda,
944945
definition: Definition<'db>,
945946
) {
946947
let ast::ParameterWithDefault {
@@ -950,12 +951,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
950951
node_index: _,
951952
} = parameter_with_default;
952953

953-
let db = self.db();
954-
955954
let default_expr = default.as_ref();
956-
let ty = if let Some(default_expr) = default_expr {
955+
let ty = if let Some(parameter_type) = self.annotated_lambda_parameter_type(index, lambda) {
956+
parameter_type
957+
} else if let Some(default_expr) = default_expr {
957958
let default_ty = self.file_expression_type(default_expr);
958-
UnionType::from_two_elements(db, Type::unknown(), default_ty)
959+
UnionType::from_two_elements(self.db(), Type::unknown(), default_ty)
959960
} else {
960961
Type::unknown()
961962
};
@@ -968,31 +969,56 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
968969
/// in a lambda expression.
969970
pub(super) fn infer_variadic_positional_lambda_parameter_definition(
970971
&mut self,
971-
_index: usize,
972+
index: usize,
972973
parameter: &'ast ast::Parameter,
974+
lambda: &'ast ast::ExprLambda,
973975
definition: Definition<'db>,
974976
) {
975-
let db = self.db();
976-
977-
let inferred_ty = Type::homogeneous_tuple(db, Type::unknown());
977+
// Note that this currently always returns `None` because we do not support `Unpack`
978+
// annotations for callable types.
979+
let ty = if let Some(parameter_type) = self.annotated_lambda_parameter_type(index, lambda) {
980+
parameter_type
981+
} else {
982+
Type::homogeneous_tuple(self.db(), Type::unknown())
983+
};
978984
self.add_binding(parameter.into(), definition)
979-
.insert(self, inferred_ty);
985+
.insert(self, ty);
980986
}
981987

982988
/// Set initial declared/inferred types for a `**kwargs` keyword-variadic parameter
983989
/// in a lambda expression.
984990
pub(super) fn infer_variadic_keyword_lambda_parameter_definition(
985991
&mut self,
986-
_index: usize,
987992
parameter: &'ast ast::Parameter,
988993
definition: Definition<'db>,
989994
) {
990-
let db = self.db();
991-
992-
let inferred_ty = KnownClass::Dict
993-
.to_specialized_instance(db, &[KnownClass::Str.to_instance(db), Type::unknown()]);
995+
let inferred_ty = KnownClass::Dict.to_specialized_instance(
996+
self.db(),
997+
&[KnownClass::Str.to_instance(self.db()), Type::unknown()],
998+
);
994999

9951000
self.add_binding(parameter.into(), definition)
9961001
.insert(self, inferred_ty);
9971002
}
1003+
1004+
/// Returns the annotated type of the lambda parameter at the given index in the provided
1005+
/// lambda expression, based on a `Callable` type annotation, if present.
1006+
fn annotated_lambda_parameter_type(
1007+
&mut self,
1008+
index: usize,
1009+
lambda: &'ast ast::ExprLambda,
1010+
) -> Option<Type<'db>> {
1011+
let scope_inference = infer_complete_scope_types(self.db(), self.scope());
1012+
let callable = scope_inference.expression_type(lambda).as_callable()?;
1013+
let [signature] = callable.signatures(self.db()).overloads.as_slice() else {
1014+
panic!("lambda callable types cannot be overloaded");
1015+
};
1016+
1017+
let parameter_type = signature.parameters().as_slice()[index].annotated_type();
1018+
if parameter_type.is_unknown() || parameter_type.has_unspecialized_type_var(self.db()) {
1019+
None
1020+
} else {
1021+
Some(parameter_type)
1022+
}
1023+
}
9981024
}

0 commit comments

Comments
 (0)
⚡