[ty] Improve invalid-argument-type diagnostics where a union type was provided (#21044)

crates/ty_python_semantic/resources/mdtest/call/function.md

@@ -820,6 +820,30 @@ def f(x: int = 1, y: str = "foo") -> int:
 reveal_type(f(y=2, x="bar"))  # revealed: int
 ```
 
+### Diagnostics for union types where the union is not assignable
+
+<!-- snapshot-diagnostics -->
+
+```py
+from typing import Sized
+
+class Foo: ...
+class Bar: ...
+class Baz: ...
+
+def f(x: Sized): ...
+def g(
+    a: str | Foo,
+    b: list[str] | str | dict[str, str] | tuple[str, ...] | bytes | frozenset[str] | set[str] | Foo,
+    c: list[str] | str | dict[str, str] | tuple[str, ...] | bytes | frozenset[str] | set[str] | Foo | Bar,
+    d: list[str] | str | dict[str, str] | tuple[str, ...] | bytes | frozenset[str] | set[str] | Foo | Bar | Baz,
+):
+    f(a)  # error: [invalid-argument-type]
+    f(b)  # error: [invalid-argument-type]
+    f(c)  # error: [invalid-argument-type]
+    f(d)  # error: [invalid-argument-type]
+```
+
 ## Too many positional arguments
 
 ### One too many

crates/ty_python_semantic/resources/mdtest/snapshots/function.md_-_Call_expression_-_Wrong_argument_type_-_Diagnostics_for_unio…_(5396a8f9e7f88f71).snap

@@ -0,0 +1,135 @@
+---
+source: crates/ty_test/src/lib.rs
+expression: snapshot
+---
+---
+mdtest name: function.md - Call expression - Wrong argument type - Diagnostics for union types where the union is not assignable
+mdtest path: crates/ty_python_semantic/resources/mdtest/call/function.md
+---
+
+# Python source files
+
+## mdtest_snippet.py
+
+```
+ 1 | from typing import Sized
+ 2 | 
+ 3 | class Foo: ...
+ 4 | class Bar: ...
+ 5 | class Baz: ...
+ 6 | 
+ 7 | def f(x: Sized): ...
+ 8 | def g(
+ 9 |     a: str | Foo,
+10 |     b: list[str] | str | dict[str, str] | tuple[str, ...] | bytes | frozenset[str] | set[str] | Foo,
+11 |     c: list[str] | str | dict[str, str] | tuple[str, ...] | bytes | frozenset[str] | set[str] | Foo | Bar,
+12 |     d: list[str] | str | dict[str, str] | tuple[str, ...] | bytes | frozenset[str] | set[str] | Foo | Bar | Baz,
+13 | ):
+14 |     f(a)  # error: [invalid-argument-type]
+15 |     f(b)  # error: [invalid-argument-type]
+16 |     f(c)  # error: [invalid-argument-type]
+17 |     f(d)  # error: [invalid-argument-type]
+```
+
+# Diagnostics
+
+```
+error[invalid-argument-type]: Argument to function `f` is incorrect
+  --> src/mdtest_snippet.py:14:7
+   |
+12 |     d: list[str] | str | dict[str, str] | tuple[str, ...] | bytes | frozenset[str] | set[str] | Foo | Bar | Baz,
+13 | ):
+14 |     f(a)  # error: [invalid-argument-type]
+   |       ^ Expected `Sized`, found `str | Foo`
+15 |     f(b)  # error: [invalid-argument-type]
+16 |     f(c)  # error: [invalid-argument-type]
+   |
+info: Element `Foo` of this union is not assignable to `Sized`
+info: Function defined here
+ --> src/mdtest_snippet.py:7:5
+  |
+5 | class Baz: ...
+6 |
+7 | def f(x: Sized): ...
+  |     ^ -------- Parameter declared here
+8 | def g(
+9 |     a: str | Foo,
+  |
+info: rule `invalid-argument-type` is enabled by default
+
+```
+
+```
+error[invalid-argument-type]: Argument to function `f` is incorrect
+  --> src/mdtest_snippet.py:15:7
+   |
+13 | ):
+14 |     f(a)  # error: [invalid-argument-type]
+15 |     f(b)  # error: [invalid-argument-type]
+   |       ^ Expected `Sized`, found `list[str] | str | dict[str, str] | ... omitted 5 union elements`
+16 |     f(c)  # error: [invalid-argument-type]
+17 |     f(d)  # error: [invalid-argument-type]
+   |
+info: Element `Foo` of this union is not assignable to `Sized`
+info: Function defined here
+ --> src/mdtest_snippet.py:7:5
+  |
+5 | class Baz: ...
+6 |
+7 | def f(x: Sized): ...
+  |     ^ -------- Parameter declared here
+8 | def g(
+9 |     a: str | Foo,
+  |
+info: rule `invalid-argument-type` is enabled by default
+
+```
+
+```
+error[invalid-argument-type]: Argument to function `f` is incorrect
+  --> src/mdtest_snippet.py:16:7
+   |
+14 |     f(a)  # error: [invalid-argument-type]
+15 |     f(b)  # error: [invalid-argument-type]
+16 |     f(c)  # error: [invalid-argument-type]
+   |       ^ Expected `Sized`, found `list[str] | str | dict[str, str] | ... omitted 6 union elements`
+17 |     f(d)  # error: [invalid-argument-type]
+   |
+info: Union elements `Foo` and `Bar` are not assignable to `Sized`
+info: Function defined here
+ --> src/mdtest_snippet.py:7:5
+  |
+5 | class Baz: ...
+6 |
+7 | def f(x: Sized): ...
+  |     ^ -------- Parameter declared here
+8 | def g(
+9 |     a: str | Foo,
+  |
+info: rule `invalid-argument-type` is enabled by default
+
+```
+
+```
+error[invalid-argument-type]: Argument to function `f` is incorrect
+  --> src/mdtest_snippet.py:17:7
+   |
+15 |     f(b)  # error: [invalid-argument-type]
+16 |     f(c)  # error: [invalid-argument-type]
+17 |     f(d)  # error: [invalid-argument-type]
+   |       ^ Expected `Sized`, found `list[str] | str | dict[str, str] | ... omitted 7 union elements`
+   |
+info: Union element `Foo`, and 2 more union elements, are not assignable to `Sized`
+info: Function defined here
+ --> src/mdtest_snippet.py:7:5
+  |
+5 | class Baz: ...
+6 |
+7 | def f(x: Sized): ...
+  |     ^ -------- Parameter declared here
+8 | def g(
+9 |     a: str | Foo,
+  |
+info: rule `invalid-argument-type` is enabled by default
+
+```

crates/ty_python_semantic/src/types/call/bind.rs

@@ -3519,6 +3519,36 @@ impl<'db> BindingError<'db> {
                     "Expected `{expected_ty_display}`, found `{provided_ty_display}`"
                 ));
 
+                if let Type::Union(union) = provided_ty {
+                    let union_elements = union.elements(context.db());
+                    let invalid_elements: Vec<Type<'db>> = union
+                        .elements(context.db())
+                        .iter()
+                        .filter(|element| !element.is_assignable_to(context.db(), *expected_ty))
+                        .copied()
+                        .collect();
+                    let first_invalid_element = invalid_elements[0].display(context.db());
+                    if invalid_elements.len() < union_elements.len() {
+                        match &invalid_elements[1..] {
+                            [] => diag.info(format_args!(
+                                "Element `{first_invalid_element}` of this union \
+                                is not assignable to `{expected_ty_display}`",
+                            )),
+                            [single] => diag.info(format_args!(
+                                "Union elements `{first_invalid_element}` and `{}` \
+                                are not assignable to `{expected_ty_display}`",
+                                single.display(context.db()),
+                            )),
+                            rest => diag.info(format_args!(
+                                "Union element `{first_invalid_element}`, \
+                                and {} more union elements, \
+                                are not assignable to `{expected_ty_display}`",
+                                rest.len(),
+                            )),
+                        }
+                    }
+                }
+
                 if let Some(matching_overload) = matching_overload {
                     if let Some((name_span, parameter_span)) =
                         matching_overload.get(context.db()).and_then(|overload| {