Repository: dask/dask
Pull request #10362 (2023-06-20 10:02:56)

mypy support for numpy 1.25
mypy has started failing (not in CI yet, since CI is still using a cache with numpy 1.24):
```
dask/array/routines.py:1569: error: Module has no attribute "round_"  [attr-defined]
Found 1 error in 1 file (checked 255 source files)
```


Patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
index b967cda3b..a05493560 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -1566,7 +1566,7 @@ def corrcoef(x, y=None, rowvar=1):
     return (c / sqr_d) / sqr_d.T
 
 
-@implements(np.round, np.round_)
+@implements(np.round)
 @derived_from(np)
 def round(a, decimals=0):
     return a.map_blocks(np.round, decimals=decimals, dtype=a.dtype)
diff --git a/dask/base.py b/dask/base.py
index c9e0491b7..32d23b8b3 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import atexit
 import dataclasses
 import datetime
 import hashlib
@@ -7,6 +8,7 @@ import inspect
 import os
 import pathlib
 import pickle
+import sys
 import threading
 import uuid
 import warnings
@@ -15,7 +17,7 @@ from collections.abc import Callable, Iterator, Mapping
 from concurrent.futures import Executor
 from contextlib import contextmanager
 from enum import Enum
-from functools import partial
+from functools import partial, wraps
 from numbers import Integral, Number
 from operator import getitem
 from typing import TYPE_CHECKING, Literal, Protocol
@@ -31,7 +33,14 @@ from dask.core import literal, quote
 from dask.hashing import hash_buffer_hex
 from dask.system import CPU_COUNT
 from dask.typing import SchedulerGetCallable
-from dask.utils import Dispatch, apply, ensure_dict, is_namedtuple_instance, key_split
+from dask.utils import (
+    Dispatch,
+    apply,
+    ensure_dict,
+    is_namedtuple_instance,
+    key_split,
+    shorten_traceback,
+)
 
 __all__ = (
     "DaskMethodsMixin",
@@ -53,6 +62,43 @@ if TYPE_CHECKING:
     from _typeshed import ReadableBuffer
 
 
+def _clean_traceback_hook(func):
+    @wraps(func)
+    def wrapper(exc_type, exc, tb):
+        tb = shorten_traceback(tb)
+        return func(exc_type, exc.with_traceback(tb), tb)
+
+    return wrapper
+
+
+def _clean_ipython_traceback(self, etype, value, tb, tb_offset=None):
+    short_tb = shorten_traceback(tb)
+    short_exc = value.with_traceback(short_tb)
+    stb = self.InteractiveTB.structured_traceback(
+        etype, short_exc, short_tb, tb_offset=tb_offset
+    )
+    self._showtraceback(type, short_exc, stb)
+
+
+try:
+    from IPython import get_ipython
+except ImportError:
+    pass
+else:
+    # if we're running in ipython, customize exception handling
+    ip = get_ipython()
+    if ip is not None:
+        ip.set_custom_exc((Exception,), _clean_ipython_traceback)
+
+
+def _restore_excepthook():
+    sys.excepthook = original_excepthook
+
+
+original_excepthook = sys.excepthook
+sys.excepthook = _clean_traceback_hook(sys.excepthook)
+
+
 @contextmanager
 def annotate(**annotations):
     """Context Manager for setting HighLevelGraph Layer annotations.
@@ -1532,3 +1578,6 @@ def clone_key(key, seed):
         prefix = key_split(key)
         return prefix + "-" + tokenize(key, seed)
     raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
+
+
+atexit.register(_restore_excepthook)
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
index 5aa9f6c38..d481de03a 100644
--- a/dask/dask-schema.yaml
+++ b/dask/dask-schema.yaml
@@ -220,3 +220,30 @@ properties:
               comprehensible, but it comes at the cost of additional processing. If
               False, then the top-most key will be used. For advanced usage, a function
               to create the new name is also accepted.
+
+  admin:
+    type: object
+    properties:
+
+      traceback:
+        type: object
+        properties:
+
+          shorten:
+            type: object
+            description:
+              Clean up tracebacks for readability.
+
+              If any of the modules in a traceback matches a regular expression listed
+              in 'when', then remove from it all modules that match a regular expression
+              in 'what'. Always preserve the first and last module.
+            additionalProperties: false
+            properties:
+              when:
+                type: array
+                items:
+                  type: string
+              what:
+                type: array
+                items:
+                  type: string
diff --git a/dask/dask.yaml b/dask/dask.yaml
index 1c6d1c919..ef5d6a505 100644
--- a/dask/dask.yaml
+++ b/dask/dask.yaml
@@ -38,3 +38,24 @@ optimization:
     max-depth-new-edges: null  # ave_width * 1.5
     subgraphs: null  # true for dask.dataframe, false for everything else
     rename-keys: true
+
+admin:
+  traceback:
+    shorten:
+      when:
+        - dask[\\\/]base.py
+        - distributed[\\\/]client.py
+      what:
+        - dask[\\\/]base.py
+        - dask[\\\/]core.py
+        - dask[\\\/]array[\\\/]core.py
+        - dask[\\\/]optimization.py
+        - dask[\\\/]dataframe[\\\/]core.py
+        - dask[\\\/]dataframe[\\\/]methods.py
+        - dask[\\\/]utils.py
+        - distributed[\\\/]worker.py
+        - distributed[\\\/]scheduler.py
+        - distributed[\\\/]client.py
+        - distributed[\\\/]utils.py
+        - tornado[\\\/]gen.py
+        - pandas[\\\/]core[\\\/]
diff --git a/dask/utils.py b/dask/utils.py
index aca551306..ed9c70d8d 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -8,6 +8,7 @@ import re
 import shutil
 import sys
 import tempfile
+import traceback
 import uuid
 import warnings
 from collections.abc import Hashable, Iterable, Iterator, Mapping, Set
@@ -2121,3 +2122,54 @@ def get_meta_library(like):
         like = like._meta
 
     return import_module(typename(like).partition(".")[0])
+
+
+def shorten_traceback(exc_traceback):
+    """Remove irrelevant stack elements from traceback.
+
+    * only shortens traceback if any of the traceback lines match
+      `admin.traceback.shorten.when`
+    * omits frames from modules that match `admin.traceback.shorten.what`
+    * always keeps the first and last frame.
+
+    Parameters
+    ----------
+    exc_traceback : types.TracebackType
+        Original traceback
+
+    Returns
+    -------
+    types.TracebackType
+        Shortened traceback
+    """
+    when_paths = config.get("admin.traceback.shorten.when")
+    what_paths = config.get("admin.traceback.shorten.what")
+    if not when_paths or not what_paths:
+        return exc_traceback
+
+    when_exp = re.compile(".*(" + "|".join(when_paths) + ")")
+    for f, _ in traceback.walk_tb(exc_traceback):
+        if when_exp.match(f.f_code.co_filename):
+            break
+    else:
+        return exc_traceback
+
+    what_exp = re.compile(".*(" + "|".join(what_paths) + ")")
+    curr = exc_traceback
+    prev = None
+
+    while curr:
+        if prev is None:
+            # always keep first frame
+            prev = curr
+        elif not curr.tb_next:
+            # always keep last frame
+            prev.tb_next = curr
+            prev = prev.tb_next
+        elif not what_exp.match(curr.tb_frame.f_code.co_filename):
+            # keep if module is not listed in what
+            prev.tb_next = curr
+            prev = prev.tb_next
+        curr = curr.tb_next
+
+    return exc_traceback


Test patch:
diff --git a/dask/tests/test_cli.py b/dask/tests/test_cli.py
index a6c2840f7..d22a46e11 100644
--- a/dask/tests/test_cli.py
+++ b/dask/tests/test_cli.py
@@ -38,7 +38,7 @@ def test_config_list():
     runner = CliRunner()
     result = runner.invoke(dask.cli.config_list)
     assert result.exit_code == 0
-    assert result.output.startswith("array:")
+    assert "array:" in result.output
 
 
 def test_version():
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index befd50762..6ff112fd8 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -6,6 +6,7 @@ distributed = pytest.importorskip("distributed")
 
 import asyncio
 import os
+import subprocess
 import sys
 from functools import partial
 from operator import add
@@ -19,6 +20,7 @@ from distributed.utils_test import (  # noqa F401
     gen_cluster,
     loop,
     loop_in_thread,
+    popen,
     varying,
 )
 
@@ -1034,3 +1036,71 @@ async def test_bag_groupby_default(c, s, a, b):
     b = db.range(100, npartitions=10)
     b2 = b.groupby(lambda x: x % 13)
     assert not any("partd" in k[0] for k in b2.dask)
+
+
+def test_shorten_traceback_excepthook(tmp_path):
+    """
+    See Also
+    --------
+    test_distributed.py::test_shorten_traceback_ipython
+    test_utils.py::test_shorten_traceback
+    """
+    client_script = """
+from dask.distributed import Client
+if __name__ == "__main__":
+    f1 = lambda: 2 / 0
+    f2 = lambda: f1() + 5
+    f3 = lambda: f2() + 1
+    with Client() as client:
+        client.submit(f3).result()
+    """
+    with open(tmp_path / "script.py", mode="w") as f:
+        f.write(client_script)
+
+    proc_args = [sys.executable, os.path.join(tmp_path, "script.py")]
+    with popen(proc_args, capture_output=True) as proc:
+        out, err = proc.communicate(timeout=60)
+
+    lines = out.decode("utf-8").split("\n")
+    lines = [line for line in lines if line.startswith("  File ")]
+
+    assert len(lines) == 4
+    assert 'script.py", line 8, in <module>' in lines[0]
+    assert 'script.py", line 6, in <lambda>' in lines[1]
+    assert 'script.py", line 5, in <lambda>' in lines[2]
+    assert 'script.py", line 4, in <lambda>' in lines[3]
+
+
+def test_shorten_traceback_ipython(tmp_path):
+    """
+    See Also
+    --------
+    test_distributed.py::test_shorten_traceback_excepthook
+    test_utils.py::test_shorten_traceback
+    """
+    pytest.importorskip("IPython", reason="Requires IPython")
+
+    client_script = """
+from dask.distributed import Client
+f1 = lambda: 2 / 0
+f2 = lambda: f1() + 5
+f3 = lambda: f2() + 1
+with Client() as client: client.submit(f3).result()
+"""
+    with popen(["ipython"], capture_output=True, stdin=subprocess.PIPE) as proc:
+        out, err = proc.communicate(input=client_script.encode(), timeout=60)
+
+    lines = out.decode("utf-8").split("\n")
+    lines = [
+        line
+        for line in lines
+        if line.startswith("File ")
+        or line.startswith("Cell ")
+        or "<ipython-input" in line
+    ]
+
+    assert len(lines) == 4
+    assert "In[5]" in lines[0] or "<ipython-input-5-" in lines[0]
+    assert "In[4]" in lines[1] or "<ipython-input-4-" in lines[1]
+    assert "In[3]" in lines[2] or "<ipython-input-3-" in lines[2]
+    assert "In[2]" in lines[3] or "<ipython-input-2-" in lines[3]
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
index fef91904f..358bc9ba9 100644
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -4,12 +4,13 @@ import datetime
 import functools
 import operator
 import pickle
+import traceback
 from array import array
 
 import pytest
 from tlz import curry
 
-from dask import get
+from dask import config, get
 from dask.highlevelgraph import HighLevelGraph
 from dask.optimization import SubgraphCallable
 from dask.utils import (
@@ -41,6 +42,7 @@ from dask.utils import (
     parse_timedelta,
     partial_by_order,
     random_state_data,
+    shorten_traceback,
     skip_doctest,
     stringify,
     stringify_collection_keys,
@@ -935,3 +937,49 @@ def test_get_meta_library_gpu():
     assert get_meta_library(cp.ndarray([])) == get_meta_library(
         da.from_array([]).to_backend("cupy")
     )
+
+
+@pytest.mark.parametrize(
+    "when,what,expect",
+    [
+        ([], [], 4),
+        ([".*"], [], 4),
+        ([], [".*"], 4),
+        ([r"nomatch"], [".*"], 4),
+        ([r".*"], ["nomatch"], 4),
+        ([".*"], [".*"], 2),
+        ([r"dask[\\\/]tests"], [], 4),
+        ([r"dask[\\\/]tests"], [r"dask[\\\/]tests"], 2),
+        ([], [r"dask[\\\/]tests"], 4),
+    ],
+)
+def test_shorten_traceback(when, what, expect):
+    """
+    See also
+    --------
+    test_distributed.py::test_shorten_traceback_excepthook
+    test_distributed.py::test_shorten_traceback_ipython
+    """
+
+    def f1():
+        return 2 / 0
+
+    def f2():
+        return f1() + 5
+
+    def f3():
+        return f2() + 1
+
+    with pytest.raises(ZeroDivisionError) as ex:
+        f3()
+
+    tb = ex.value.__traceback__
+    with config.set(
+        {
+            "admin.traceback.shorten.when": when,
+            "admin.traceback.shorten.what": what,
+        }
+    ):
+        tb = shorten_traceback(tb)
+    frame_count = len(list(traceback.walk_tb(tb)))
+    assert frame_count == expect