fixes #1631

Author: jph00

nbdev/_modidx.py

@@ -378,7 +378,8 @@
                             'nbdev.sync._update_nb': ('api/sync.html#_update_nb', 'nbdev/sync.py'),
                             'nbdev.sync.absolute_import': ('api/sync.html#absolute_import', 'nbdev/sync.py'),
                             'nbdev.sync.nbdev_update': ('api/sync.html#nbdev_update', 'nbdev/sync.py')},
-            'nbdev.test': { 'nbdev.test._int_handler': ('api/test.html#_int_handler', 'nbdev/test.py'),
+            'nbdev.test': { 'nbdev.test._await_chain': ('api/test.html#_await_chain', 'nbdev/test.py'),
+                            'nbdev.test._int_handler': ('api/test.html#_int_handler', 'nbdev/test.py'),
                             'nbdev.test._keep_file': ('api/test.html#_keep_file', 'nbdev/test.py'),
                             'nbdev.test._test_nb_sync': ('api/test.html#_test_nb_sync', 'nbdev/test.py'),
                             'nbdev.test.nbdev_test': ('api/test.html#nbdev_test', 'nbdev/test.py'),

nbdev/test.py

@@ -8,7 +8,7 @@
 __all__ = ['test_nb', 'nbdev_test']
 
 # %% ../nbs/api/12_test.ipynb #45e10c3f
-import time,os,sys,traceback,contextlib,inspect,signal,asyncio
+import time,os,sys,io,traceback,contextlib,inspect,signal,asyncio
 from fastcore.basics import *
 from fastcore.imports import *
 from fastcore.foundation import *
@@ -27,11 +27,24 @@
 # %% ../nbs/api/12_test.ipynb #dc8994ac
 _cur_nb = [None]
 
+def _await_chain(t):
+    "One frame per coroutine in task `t`'s await chain, deepest last: where a suspended hang actually sits"
+    co = t.get_coro()
+    while co is not None:
+        f = getattr(co, 'cr_frame', None) or getattr(co, 'ag_frame', None) or getattr(co, 'gi_frame', None)
+        if f is not None: yield f, f.f_lineno
+        co = getattr(co, 'cr_await', None) or getattr(co, 'ag_await', None) or getattr(co, 'gi_yieldfrom', None)
+
 def _int_handler(signum, frame):
     "Dump the running notebook's stack and exit; installed on SIGINT by `test_nb`"
     if _cur_nb[0] is not None:
-        stk = ''.join(traceback.format_stack(frame))
-        os.write(2, f'\n=== nbdev-test interrupted: {_cur_nb[0]} ===\n{stk}'.encode())
+        buf = io.StringIO()
+        traceback.print_stack(frame, file=buf)
+        with contextlib.suppress(RuntimeError):  # no running loop: sync frames already cover it
+            for t in asyncio.all_tasks():
+                print(f'\n{t}', file=buf)
+                buf.writelines(traceback.StackSummary.extract(_await_chain(t)).format())
+        os.write(2, f'\n=== nbdev-test interrupted: {_cur_nb[0]} ===\n{buf.getvalue()}'.encode())
     os._exit(130)
 
 # %% ../nbs/api/12_test.ipynb #3f4fa1ad

nbs/api/12_test.ipynb

@@ -28,7 +28,7 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "import time,os,sys,traceback,contextlib,inspect,signal,asyncio\n",
+    "import time,os,sys,io,traceback,contextlib,inspect,signal,asyncio\n",
     "from fastcore.basics import *\n",
     "from fastcore.imports import *\n",
     "from fastcore.foundation import *\n",
@@ -50,7 +50,7 @@
    "id": "7fe8a875",
    "metadata": {},
    "source": [
-    "When a test run hangs, ctrl-c should say what each notebook was executing. Each worker installs a SIGINT handler that prints the in-flight notebook's stack as a single block (one `os.write`, so parallel workers' dumps don't interleave) and then exits immediately, rather than letting the interrupt surface inside the cell where the shell would catch it and carry on. An idle worker exits silently."
+    "When a test run hangs, ctrl-c should say what each notebook was executing. Each worker installs a SIGINT handler that prints the in-flight notebook's stack — the sync frames plus every pending asyncio task, since a hang inside an awaited coroutine lives on the task, not the sync stack — as a single block (one `os.write`, so parallel workers' dumps don't interleave) and then exits immediately, rather than letting the interrupt surface inside the cell where the shell would catch it and carry on. An idle worker exits silently."
    ]
   },
   {
@@ -63,11 +63,24 @@
     "#| export\n",
     "_cur_nb = [None]\n",
     "\n",
+    "def _await_chain(t):\n",
+    "    \"One frame per coroutine in task `t`'s await chain, deepest last: where a suspended hang actually sits\"\n",
+    "    co = t.get_coro()\n",
+    "    while co is not None:\n",
+    "        f = getattr(co, 'cr_frame', None) or getattr(co, 'ag_frame', None) or getattr(co, 'gi_frame', None)\n",
+    "        if f is not None: yield f, f.f_lineno\n",
+    "        co = getattr(co, 'cr_await', None) or getattr(co, 'ag_await', None) or getattr(co, 'gi_yieldfrom', None)\n",
+    "\n",
     "def _int_handler(signum, frame):\n",
     "    \"Dump the running notebook's stack and exit; installed on SIGINT by `test_nb`\"\n",
     "    if _cur_nb[0] is not None:\n",
-    "        stk = ''.join(traceback.format_stack(frame))\n",
-    "        os.write(2, f'\\n=== nbdev-test interrupted: {_cur_nb[0]} ===\\n{stk}'.encode())\n",
+    "        buf = io.StringIO()\n",
+    "        traceback.print_stack(frame, file=buf)\n",
+    "        with contextlib.suppress(RuntimeError):  # no running loop: sync frames already cover it\n",
+    "            for t in asyncio.all_tasks():\n",
+    "                print(f'\\n{t}', file=buf)\n",
+    "                buf.writelines(traceback.StackSummary.extract(_await_chain(t)).format())\n",
+    "        os.write(2, f'\\n=== nbdev-test interrupted: {_cur_nb[0]} ===\\n{buf.getvalue()}'.encode())\n",
     "    os._exit(130)"
    ]
   },