Repository: dask/dask
Pull request: Improve exception for `validate_key`
- [x] Closes #10764
- [ ] Tests added / passed
- [x] Passes `pre-commit run --all-files`

The exception now indicates the index of the invalid key type and the traceback reveals the actual problem. The latter is important for nested tuples.

```
Traceback (most recent call last):
  File "/Users/hendrikmakait/projects/dask/dask/dask/core.py", line 235, in validate_key
    validate_key(part)
  File "/Users/hendrikmakait/projects/dask/dask/dask/core.py", line 240, in validate_key
    raise TypeError(f"Unexpected key type {typ} ({key=!r})")
TypeError: Unexpected key type <class 'numpy.int64'> (key=2)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/hendrikmakait/projects/dask/dask/dask/core.py", line 237, in validate_key
    raise TypeError(
TypeError: Composite key contains unexpected key type at _index=1 (key=(1, 2, 3))
```

Problem statement:
Exception from `validate_key` is misleading for composite keys
When validating a key tuple with an invalid component, the exception is currently not helpful for identifying the problem and might even be misleading:

```python
import numpy as np
from dask.core import validate_key
validate_key((1, np.int64(2), 3))
```
creates the following traceback:

```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/hendrikmakait/projects/dask/dask/dask/core.py", line 223, in validate_key
    raise TypeError(f"Unexpected key type {type(key)} (value: {key!r})")
TypeError: Unexpected key type <class 'tuple'> (value: (1, 2, 3))
```

That is, the exception suggests that `tuple` is an invalid type a key even though it is the _only_ valid type for creating composite keys. The printed key also doesn't give any hint that the `2` is to blame here.

Hints from issue discussion:

Commit message: Improve exception

Patch to review:
diff --git a/dask/core.py b/dask/core.py
index ef75ca563cd..4121a034c38 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -218,9 +218,26 @@ def iskey(key: object) -> bool:
 
 
 def validate_key(key: object) -> None:
-    """Validate the format of a dask key."""
-    if not iskey(key):
-        raise TypeError(f"Unexpected key type {type(key)} (value: {key!r})")
+    """Validate the format of a dask key.
+
+    See Also
+    --------
+    iskey
+    """
+    if iskey(key):
+        return
+    typ = type(key)
+
+    if typ is tuple:
+        _index = None
+        try:
+            for _index, part in enumerate(cast(tuple, key)):
+                validate_key(part)
+        except TypeError as e:
+            raise TypeError(
+                f"Composite key contains unexpected key type at {_index=} ({key=!r})"
+            ) from e
+    raise TypeError(f"Unexpected key type {typ} ({key=!r})")
 
 
 @overload
diff --git a/dask/tests/test_core.py b/dask/tests/test_core.py
index a3d72380dd7..41b42b44d35 100644
--- a/dask/tests/test_core.py
+++ b/dask/tests/test_core.py
@@ -73,9 +73,12 @@ def test_iskey_numpy_types():
 def test_validate_key():
     validate_key(1)
     validate_key(("x", 1))
-    with pytest.raises(TypeError):
+    with pytest.raises(TypeError, match="Unexpected key type.*list"):
         validate_key(["x", 1])
 
+    with pytest.raises(TypeError, match="unexpected key type at index=1"):
+        validate_key((2, int))
+
 
 def test_istask():
     assert istask((inc, 1))


Code review comments:
diff_hunk: @@ -218,9 +218,26 @@ def iskey(key: object) -> bool:
 
 
 def validate_key(key: object) -> None:
-    """Validate the format of a dask key."""
-    if not iskey(key):
-        raise TypeError(f"Unexpected key type {type(key)} (value: {key!r})")
+    """Validate the format of a dask key.
+
+    See Also
+    --------
+    iskey
+    """
+    if iskey(key):
+        return
+    typ = type(key)
+
+    if typ is tuple:
+        _index = None
+        try:
+            for _index, part in enumerate(cast(tuple, key)):
+                validate_key(part)
+        except TypeError as e:
+            raise TypeError(
+                f"Composite key contains unexpected key type at {_index=} ({key=!r})"
path: dask/core.py
text: @user1:
```suggestion
                f"Composite key contains unexpected key type at index={_index} ({key=!r})"
```

That would get ci passing

@author:
I shouldn't have followed `flake8`s advice there. Reverted to `index` and ignoring the check.

@user1:
Sounds good to me