[ty] Function argument inlay hints (#19269)

crates/ty_ide/src/inlay_hints.rs

@@ -6,7 +6,7 @@ use ruff_python_ast::{AnyNodeRef, Expr, Stmt};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 use std::fmt;
 use std::fmt::Formatter;
-use ty_python_semantic::types::Type;
+use ty_python_semantic::types::{Type, inlay_hint_function_argument_details};
 use ty_python_semantic::{HasType, SemanticModel};
 
 #[derive(Debug, Clone, Eq, PartialEq)]
@@ -24,7 +24,7 @@ impl<'db> InlayHint<'db> {
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum InlayHintContent<'db> {
     Type(Type<'db>),
-    ReturnType(Type<'db>),
+    FunctionArgumentName(String),
 }
 
 impl<'db> InlayHintContent<'db> {
@@ -44,8 +44,8 @@ impl fmt::Display for DisplayInlayHint<'_, '_> {
             InlayHintContent::Type(ty) => {
                 write!(f, ": {}", ty.display(self.db))
             }
-            InlayHintContent::ReturnType(ty) => {
-                write!(f, " -> {}", ty.display(self.db))
+            InlayHintContent::FunctionArgumentName(name) => {
+                write!(f, "{name}=")
             }
         }
     }
@@ -76,9 +76,18 @@ pub struct InlayHintSettings {
     /// x": Literal[1]" = 1
     /// ```
     pub variable_types: bool,
+    /// Whether to show function argument names.
+    ///
+    /// For example, this would enable / disable hints like the ones quoted below:
+    /// ```python
+    /// def foo(x: int): pass
+    /// foo("x="1)
+    /// ```
+    pub function_argument_names: bool,
 }
 
 struct InlayHintVisitor<'a, 'db> {
+    db: &'db dyn Db,
     model: SemanticModel<'db>,
     hints: Vec<InlayHint<'db>>,
     in_assignment: bool,
@@ -89,6 +98,7 @@ struct InlayHintVisitor<'a, 'db> {
 impl<'a, 'db> InlayHintVisitor<'a, 'db> {
     fn new(db: &'db dyn Db, file: File, range: TextRange, settings: &'a InlayHintSettings) -> Self {
         Self {
+            db,
             model: SemanticModel::new(db, file),
             hints: Vec::new(),
             in_assignment: false,
@@ -98,11 +108,29 @@ impl<'a, 'db> InlayHintVisitor<'a, 'db> {
     }
 
     fn add_type_hint(&mut self, position: TextSize, ty: Type<'db>) {
+        if !self.settings.variable_types {
+            return;
+        }
         self.hints.push(InlayHint {
             position,
             content: InlayHintContent::Type(ty),
         });
     }
+
+    fn add_function_argument_name(&mut self, position: TextSize, name: String) {
+        if !self.settings.function_argument_names {
+            return;
+        }
+
+        if name.starts_with('_') {
+            return;
+        }
+
+        self.hints.push(InlayHint {
+            position,
+            content: InlayHintContent::FunctionArgumentName(name),
+        });
+    }
 }
 
 impl SourceOrderVisitor<'_> for InlayHintVisitor<'_, '_> {
@@ -123,25 +151,23 @@ impl SourceOrderVisitor<'_> for InlayHintVisitor<'_, '_> {
 
         match stmt {
             Stmt::Assign(assign) => {
-                if !self.settings.variable_types {
-                    return;
-                }
-
                 self.in_assignment = true;
                 for target in &assign.targets {
                     self.visit_expr(target);
                 }
                 self.in_assignment = false;
 
+                self.visit_expr(&assign.value);
+
+                return;
+            }
+            Stmt::Expr(expr) => {
+                self.visit_expr(&expr.value);
                 return;
             }
             // TODO
             Stmt::FunctionDef(_) => {}
             Stmt::For(_) => {}
-            Stmt::Expr(_) => {
-                // Don't traverse into expression statements because we don't show any hints.
-                return;
-            }
             _ => {}
         }
 
@@ -149,15 +175,32 @@ impl SourceOrderVisitor<'_> for InlayHintVisitor<'_, '_> {
     }
 
     fn visit_expr(&mut self, expr: &'_ Expr) {
-        if !self.in_assignment {
-            return;
-        }
-
         match expr {
             Expr::Name(name) => {
-                if name.ctx.is_store() {
-                    let ty = expr.inferred_type(&self.model);
-                    self.add_type_hint(expr.range().end(), ty);
+                if self.in_assignment {
+                    if name.ctx.is_store() {
+                        let ty = expr.inferred_type(&self.model);
+                        self.add_type_hint(expr.range().end(), ty);
+                    }
+                }
+                source_order::walk_expr(self, expr);
+            }
+            Expr::Call(call) => {
+                let argument_names =
+                    inlay_hint_function_argument_details(self.db, &self.model, call)
+                        .map(|details| details.argument_names)
+                        .unwrap_or_default();
+
+                self.visit_expr(&call.func);
+
+                for (index, arg_or_keyword) in call.arguments.arguments_source_order().enumerate() {
+                    if let Some(name) = argument_names.get(&index) {
+                        self.add_function_argument_name(
+                            arg_or_keyword.range().start(),
+                            name.to_string(),
+                        );
+                    }
+                    self.visit_expr(arg_or_keyword.value());
                 }
             }
             _ => {
@@ -177,6 +220,7 @@ mod tests {
         files::{File, system_path_to_file},
         source::source_text,
     };
+    use ruff_python_trivia::textwrap::dedent;
     use ruff_text_size::TextSize;
 
     use ruff_db::system::{DbWithWritableSystem, SystemPathBuf};
@@ -194,6 +238,8 @@ mod tests {
             SystemPathBuf::from("/"),
         ));
 
+        let source = dedent(source);
+
         let start = source.find(START);
         let end = source
             .find(END)
@@ -245,6 +291,7 @@ mod tests {
         fn inlay_hints(&self) -> String {
             self.inlay_hints_with_settings(&InlayHintSettings {
                 variable_types: true,
+                function_argument_names: true,
             })
         }
 
@@ -314,16 +361,529 @@ mod tests {
     }
 
     #[test]
-    fn disabled_variable_types() {
+    fn test_disabled_variable_types() {
         let test = inlay_hint_test("x = 1");
 
         assert_snapshot!(
             test.inlay_hints_with_settings(&InlayHintSettings {
                 variable_types: false,
+                ..Default::default()
             }),
             @r"
         x = 1
         "
         );
     }
+
+    #[test]
+    fn test_function_call_with_positional_or_keyword_parameter() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int): pass
+            foo(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int): pass
+        foo([x=]1)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_with_positional_only_parameter() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int, /): pass
+            foo(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int, /): pass
+        foo(1)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_with_variadic_parameter() {
+        let test = inlay_hint_test(
+            "
+            def foo(*args: int): pass
+            foo(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(*args: int): pass
+        foo(1)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_with_keyword_variadic_parameter() {
+        let test = inlay_hint_test(
+            "
+            def foo(**kwargs: int): pass
+            foo(x=1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(**kwargs: int): pass
+        foo(x=1)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_with_keyword_only_parameter() {
+        let test = inlay_hint_test(
+            "
+            def foo(*, x: int): pass
+            foo(x=1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(*, x: int): pass
+        foo(x=1)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_positional_only_and_positional_or_keyword_parameters() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int, /, y: int): pass
+            foo(1, 2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int, /, y: int): pass
+        foo(1, [y=]2)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_positional_only_and_variadic_parameters() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int, /, *args: int): pass
+            foo(1, 2, 3)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int, /, *args: int): pass
+        foo(1, 2, 3)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_positional_only_and_keyword_variadic_parameters() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int, /, **kwargs: int): pass
+            foo(1, x=2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int, /, **kwargs: int): pass
+        foo(1, x=2)
+        ");
+    }
+
+    #[test]
+    fn test_class_constructor_call_init() {
+        let test = inlay_hint_test(
+            "
+            class Foo:
+                def __init__(self, x: int): pass
+            Foo(1)
+            f = Foo(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        class Foo:
+            def __init__(self, x: int): pass
+        Foo([x=]1)
+        f[: Foo] = Foo([x=]1)
+        ");
+    }
+
+    #[test]
+    fn test_class_constructor_call_new() {
+        let test = inlay_hint_test(
+            "
+            class Foo:
+                def __new__(cls, x: int): pass
+            Foo(1)
+            f = Foo(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        class Foo:
+            def __new__(cls, x: int): pass
+        Foo([x=]1)
+        f[: Foo] = Foo([x=]1)
+        ");
+    }
+
+    #[test]
+    fn test_class_constructor_call_meta_class_call() {
+        let test = inlay_hint_test(
+            "
+            class MetaFoo:
+                def __call__(self, x: int): pass
+            class Foo(metaclass=MetaFoo):
+                pass
+            Foo(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        class MetaFoo:
+            def __call__(self, x: int): pass
+        class Foo(metaclass=MetaFoo):
+            pass
+        Foo([x=]1)
+        ");
+    }
+
+    #[test]
+    fn test_callable_call() {
+        let test = inlay_hint_test(
+            "
+            from typing import Callable
+            def foo(x: Callable[[int], int]):
+                x(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        from typing import Callable
+        def foo(x: Callable[[int], int]):
+            x(1)
+        ");
+    }
+
+    #[test]
+    fn test_instance_method_call() {
+        let test = inlay_hint_test(
+            "
+            class Foo:
+                def bar(self, y: int): pass
+            Foo().bar(2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        class Foo:
+            def bar(self, y: int): pass
+        Foo().bar([y=]2)
+        ");
+    }
+
+    #[test]
+    fn test_class_method_call() {
+        let test = inlay_hint_test(
+            "
+            class Foo:
+                @classmethod
+                def bar(cls, y: int): pass
+            Foo.bar(2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        class Foo:
+            @classmethod
+            def bar(cls, y: int): pass
+        Foo.bar([y=]2)
+        ");
+    }
+
+    #[test]
+    fn test_static_method_call() {
+        let test = inlay_hint_test(
+            "
+            class Foo:
+                @staticmethod
+                def bar(y: int): pass
+            Foo.bar(2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        class Foo:
+            @staticmethod
+            def bar(y: int): pass
+        Foo.bar([y=]2)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_with_union_type() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int | str): pass
+            foo(1)
+            foo('abc')",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int | str): pass
+        foo([x=]1)
+        foo([x=]'abc')
+        ");
+    }
+
+    #[test]
+    fn test_function_call_multiple_positional_arguments() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int, y: str, z: bool): pass
+            foo(1, 'hello', True)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int, y: str, z: bool): pass
+        foo([x=]1, [y=]'hello', [z=]True)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_mixed_positional_and_keyword() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int, y: str, z: bool): pass
+            foo(1, z=True, y='hello')",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int, y: str, z: bool): pass
+        foo([x=]1, z=True, y='hello')
+        ");
+    }
+
+    #[test]
+    fn test_function_call_with_default_parameters() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int, y: str = 'default', z: bool = False): pass
+            foo(1)
+            foo(1, 'custom')
+            foo(1, 'custom', True)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int, y: str = 'default', z: bool = False): pass
+        foo([x=]1)
+        foo([x=]1, [y=]'custom')
+        foo([x=]1, [y=]'custom', [z=]True)
+        ");
+    }
+
+    #[test]
+    fn test_nested_function_calls() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: int) -> int:
+                return x * 2
+
+            def bar(y: str) -> str:
+                return y
+
+            def baz(a: int, b: str, c: bool): pass
+
+            baz(foo(5), bar(bar('test')), True)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int) -> int:
+            return x * 2
+
+        def bar(y: str) -> str:
+            return y
+
+        def baz(a: int, b: str, c: bool): pass
+
+        baz([a=]foo([x=]5), [b=]bar([y=]bar([y=]'test')), [c=]True)
+        ");
+    }
+
+    #[test]
+    fn test_method_chaining() {
+        let test = inlay_hint_test(
+            "
+            class A:
+                def foo(self, value: int) -> 'A':
+                    return self
+                def bar(self, name: str) -> 'A':
+                    return self
+                def baz(self): pass
+            A().foo(42).bar('test').baz()",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        class A:
+            def foo(self, value: int) -> 'A':
+                return self
+            def bar(self, name: str) -> 'A':
+                return self
+            def baz(self): pass
+        A().foo([value=]42).bar([name=]'test').baz()
+        ");
+    }
+
+    #[test]
+    fn test_nexted_keyword_function_calls() {
+        let test = inlay_hint_test(
+            "
+            def foo(x: str) -> str:
+                return x
+            def bar(y: int): pass
+            bar(y=foo('test'))
+            ",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: str) -> str:
+            return x
+        def bar(y: int): pass
+        bar(y=foo([x=]'test'))
+        ");
+    }
+
+    #[test]
+    fn test_lambda_function_calls() {
+        let test = inlay_hint_test(
+            "
+            foo = lambda x: x * 2
+            bar = lambda a, b: a + b
+            foo(5)
+            bar(1, 2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        foo[: (x) -> Unknown] = lambda x: x * 2
+        bar[: (a, b) -> Unknown] = lambda a, b: a + b
+        foo([x=]5)
+        bar([a=]1, [b=]2)
+        ");
+    }
+
+    #[test]
+    fn test_complex_parameter_combinations() {
+        let test = inlay_hint_test(
+            "
+            def foo(a: int, b: str, /, c: float, d: bool = True, *, e: int, f: str = 'default'): pass
+            foo(1, 'pos', 3.14, False, e=42)
+            foo(1, 'pos', 3.14, e=42, f='custom')",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(a: int, b: str, /, c: float, d: bool = True, *, e: int, f: str = 'default'): pass
+        foo(1, 'pos', [c=]3.14, [d=]False, e=42)
+        foo(1, 'pos', [c=]3.14, e=42, f='custom')
+        ");
+    }
+
+    #[test]
+    fn test_generic_function_calls() {
+        let test = inlay_hint_test(
+            "
+            from typing import TypeVar, Generic
+
+            T = TypeVar('T')
+
+            def identity(x: T) -> T:
+                return x
+
+            identity(42)
+            identity('hello')",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r#"
+        from typing import TypeVar, Generic
+
+        T[: typing.TypeVar("T")] = TypeVar([name=]'T')
+
+        def identity(x: T) -> T:
+            return x
+
+        identity([x=]42)
+        identity([x=]'hello')
+        "#);
+    }
+
+    #[test]
+    fn test_overloaded_function_calls() {
+        let test = inlay_hint_test(
+            "
+            from typing import overload
+
+            @overload
+            def foo(x: int) -> str: ...
+            @overload
+            def foo(x: str) -> int: ...
+            def foo(x):
+                return x
+            
+            foo(42)
+            foo('hello')",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        from typing import overload
+
+        @overload
+        def foo(x: int) -> str: ...
+        @overload
+        def foo(x: str) -> int: ...
+        def foo(x):
+            return x
+
+        foo([x=]42)
+        foo([x=]'hello')
+        ");
+    }
+
+    #[test]
+    fn test_disabled_function_argument_names() {
+        let test = inlay_hint_test(
+            "
+        def foo(x: int): pass
+        foo(1)",
+        );
+
+        assert_snapshot!(test.inlay_hints_with_settings(&InlayHintSettings {
+            function_argument_names: false,
+            ..Default::default()
+        }), @r"
+        def foo(x: int): pass
+        foo(1)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_out_of_range() {
+        let test = inlay_hint_test(
+            "
+            <START>def foo(x: int): pass
+            def bar(y: int): pass
+            foo(1)<END>
+            bar(2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(x: int): pass
+        def bar(y: int): pass
+        foo([x=]1)
+        bar(2)
+        ");
+    }
+
+    #[test]
+    fn test_function_call_with_argument_name_starting_with_underscore() {
+        let test = inlay_hint_test(
+            "
+            def foo(_x: int, y: int): pass
+            foo(1, 2)",
+        );
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        def foo(_x: int, y: int): pass
+        foo(1, [y=]2)
+        ");
+    }
 }