[ty] Infer lambda parameter types with Callable type context (#24317)

Improves #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 #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.

Author: ibraheemdev

crates/ty_ide/src/semantic_tokens.rs

@@ -43,7 +43,7 @@ use ruff_python_ast::{
 };
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 use std::ops::Deref;
-use ty_python_core::definition::{Definition, DefinitionKind};
+use ty_python_core::definition::{Definition, DefinitionKind, ParameterDefinitionNodeKind};
 use ty_python_semantic::{
     HasType, SemanticModel,
     types::ide_support::{
@@ -314,7 +314,7 @@ impl<'db> SemanticTokenVisitor<'db> {
             }
             DefinitionKind::Class(_) => Some((SemanticTokenType::Class, modifiers)),
             DefinitionKind::TypeVar(_) => Some((SemanticTokenType::TypeParameter, modifiers)),
-            DefinitionKind::Parameter(parameter) => {
+            DefinitionKind::Parameter(ParameterDefinitionNodeKind::Parameter(parameter)) => {
                 let parsed = parsed_module(db, definition.file(db));
                 let ty = parameter.node(&parsed.load(db)).inferred_type(&model);
 
@@ -342,12 +342,7 @@ impl<'db> SemanticTokenVisitor<'db> {
 
                 Some((SemanticTokenType::Parameter, modifiers))
             }
-            DefinitionKind::VariadicPositionalParameter(_) => {
-                Some((SemanticTokenType::Parameter, modifiers))
-            }
-            DefinitionKind::VariadicKeywordParameter(_) => {
-                Some((SemanticTokenType::Parameter, modifiers))
-            }
+            DefinitionKind::Parameter(_) => Some((SemanticTokenType::Parameter, modifiers)),
             DefinitionKind::TypeAlias(_) => Some((SemanticTokenType::TypeParameter, modifiers)),
             DefinitionKind::Import(_)
             | DefinitionKind::ImportFrom(_)

crates/ty_python_core/src/ast_ids.rs

@@ -150,6 +150,12 @@ pub(crate) mod node_key {
         }
     }
 
+    impl From<&ast::ExprLambda> for ExpressionNodeKey {
+        fn from(value: &ast::ExprLambda) -> Self {
+            Self(NodeKey::from_node(value))
+        }
+    }
+
     impl From<&ast::Identifier> for ExpressionNodeKey {
         fn from(value: &ast::Identifier) -> Self {
             Self(NodeKey::from_node(value))