fix shorten traceback

distributed/tests/test_client.py

@@ -468,46 +468,46 @@ def test_Future_release_sync(c):
     poll_for(lambda: not c.futures, timeout=5)
 
 
-# @pytest.mark.parametrize("method", ["result", "gather"])
-# def test_short_tracebacks(c, method):
-#     """
-#     See also
-#     --------
-#     test_short_tracebacks_async
-#     dask/tests/test_traceback.py
-#     """
-#     future = c.submit(div, 1, 0)
-#     with pytest.raises(ZeroDivisionError) as e:
-#         if method == "result":
-#             future.result()
-#         else:
-#             c.gather(future)
-
-#     frames = list(traceback.walk_tb(e.value.__traceback__))
-#     assert len(frames) < 4
-
-
-# @pytest.mark.parametrize("method", ["await", "result", "gather"])
-# @gen_cluster(client=True)
-# async def test_short_tracebacks_async(c, s, a, b, method):
-#     """
-#     See also
-#     --------
-#     test_short_tracebacks
-#     dask/tests/test_traceback.py
-#     """
-#     future = c.submit(div, 1, 0)
-
-#     with pytest.raises(ZeroDivisionError) as e:
-#         if method == "await":
-#             await future
-#         elif method == "result":
-#             await future.result()
-#         else:
-#             await c.gather(future)
-
-#     frames = list(traceback.walk_tb(e.value.__traceback__))
-#     assert len(frames) < 4
+@pytest.mark.parametrize("method", ["result", "gather"])
+def test_short_tracebacks(c, method):
+    """
+    See also
+    --------
+    test_short_tracebacks_async
+    dask/tests/test_traceback.py
+    """
+    future = c.submit(div, 1, 0)
+    with pytest.raises(ZeroDivisionError) as e:
+        if method == "result":
+            future.result()
+        else:
+            c.gather(future)
+
+    frames = list(traceback.walk_tb(e.value.__traceback__))
+    assert len(frames) < 4
+
+
+@pytest.mark.parametrize("method", ["await", "result", "gather"])
+@gen_cluster(client=True)
+async def test_short_tracebacks_async(c, s, a, b, method):
+    """
+    See also
+    --------
+    test_short_tracebacks
+    dask/tests/test_traceback.py
+    """
+    future = c.submit(div, 1, 0)
+
+    with pytest.raises(ZeroDivisionError) as e:
+        if method == "await":
+            await future
+        elif method == "result":
+            await future.result()
+        else:
+            await c.gather(future)
+
+    frames = list(traceback.walk_tb(e.value.__traceback__))
+    assert len(frames) < 4
 
 
 @gen_cluster(client=True)

distributed/utils.py

@@ -952,6 +952,7 @@ def get_traceback():
         os.path.join("distributed", "scheduler"),
         os.path.join("tornado", "gen.py"),
         os.path.join("concurrent", "futures"),
+        os.path.join("dask", "task_spec"),
     ]
     while exc_traceback and any(
         b in exc_traceback.tb_frame.f_code.co_filename for b in bad