Fix Output capture from background threads

python/ipywidgets/ipywidgets/widgets/tests/test_widget_output.py

@@ -64,6 +64,41 @@ def test_set_msg_id_when_capturing(self):
                 assert widget.msg_id == msg_id
             assert widget.msg_id == ''
 
+    def test_set_parent_when_capturing(self):
+        msg_id = 'msg-id'
+        shell_parent = {'header': {'msg_id': msg_id}}
+        kernel_parent = {'header': {'msg_id': 'kernel-msg-id'}}
+        parent_calls = []
+
+        def get_kernel_parent(self_):
+            return kernel_parent
+
+        kernel = type(
+            'mock_kernel',
+            (object, ),
+            {'get_parent': get_kernel_parent}
+        )()
+
+        def get_parent(self_):
+            return shell_parent
+
+        def set_parent(self_, parent):
+            parent_calls.append(parent)
+
+        ipython = type(
+            'mock_ipython',
+            (object, ),
+            {'kernel': kernel, 'get_parent': get_parent, 'set_parent': set_parent}
+        )
+        clear_output = self._mock_clear_output()
+
+        with self._mocked_ipython(ipython, clear_output):
+            widget = widget_output.Output()
+            with widget:
+                assert widget.msg_id == msg_id
+
+        assert parent_calls == [shell_parent]
+
     def test_clear_output(self):
         msg_id = 'msg-id'
         get_ipython = self._mock_get_ipython(msg_id)

python/ipywidgets/ipywidgets/widgets/widget_output.py

@@ -115,15 +115,30 @@ def __enter__(self):
             kernel = self.comm.kernel
 
         if kernel:
-            parent = None
-            if hasattr(kernel, "get_parent"):
-                parent = kernel.get_parent()
-            elif hasattr(kernel, "_parent_header"):
-                # ipykernel < 6: kernel._parent_header is the parent *request*
-                parent = kernel._parent_header
-
-            if parent and parent.get("header"):
-                self.msg_id = parent["header"]["msg_id"]
+            parent_request = None
+            if ip and hasattr(ip, "get_parent"):
+                # Prefer the shell parent request so display hooks and streams
+                # are refreshed from the same message below.
+                shell_parent = ip.get_parent()
+                if shell_parent:
+                    parent_request = shell_parent
+            if not parent_request:
+                if hasattr(kernel, "get_parent"):
+                    # Fallback for contexts without an active shell parent:
+                    # ipykernel >= 6 keeps parent requests on the kernel.
+                    kernel_parent = kernel.get_parent()
+                elif hasattr(kernel, "_parent_header"):
+                    # ipykernel < 6: kernel._parent_header is the parent *request*
+                    kernel_parent = kernel._parent_header
+                else:
+                    kernel_parent = None
+                if kernel_parent:
+                    parent_request = kernel_parent
+
+            if parent_request and parent_request.get("header"):
+                if ip and hasattr(ip, "set_parent"):
+                    ip.set_parent(parent_request)
+                self.msg_id = parent_request["header"]["msg_id"]
                 self.__counter += 1
 
     def __exit__(self, etype, evalue, tb):

ui-tests/tests/notebooks/output-background-thread.ipynb

@@ -0,0 +1,74 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import threading\n",
+    "from time import sleep\n",
+    "\n",
+    "import ipywidgets as widgets\n",
+    "\n",
+    "out = widgets.Output()\n",
+    "\n",
+    "def write_to_output():\n",
+    "    with out:\n",
+    "        for i in range(10):\n",
+    "            print(i)\n",
+    "            sleep(0.5)\n",
+    "\n",
+    "display(out)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "threading.Thread(target=write_to_output).start()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "'first foreground cell'\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "'second foreground cell'\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

ui-tests/tests/widgets.test.ts

@@ -1,3 +1,3 @@
 // Copyright (c) Jupyter Development Team.
 // Distributed under the terms of the Modified BSD License.
 
@@ -8,7 +8,7 @@
 import * as path from 'path';
 
 test.describe('Widget Visual Regression', () => {
   test.beforeEach(async ({ page, tmpPath }) => {
     await page.contents.uploadDirectory(
       path.resolve(__dirname, './notebooks'),
       tmpPath
@@ -43,4 +43,33 @@
       expect.soft(captures[i]).toMatchSnapshot(image);
     }
   });
+
+  test('captures Output widget output from background threads', async ({
+    page,
+    tmpPath,
+  }) => {
+    const notebook = 'output-background-thread.ipynb';
+    await page.notebook.openByPath(`${tmpPath}/${notebook}`);
+    await page.notebook.activate(notebook);
+
+    const widgetOutput = page.locator('.jp-Notebook .widget-output');
+    await page.notebook.runCellByCell({
+      onAfterCellRun: async (cellIndex: number) => {
+        if (cellIndex === 1) {
+          await expect(widgetOutput).toContainText('0');
+        }
+      },
+    });
+
+    await expect(widgetOutput).toHaveText('0 1 2 3 4 5 6 7 8 9', {
+      timeout: 10_000,
+    });
+
+    expect(await page.notebook.getCellTextOutput(2)).toEqual([
+      "'first foreground cell'",
+    ]);
+    expect(await page.notebook.getCellTextOutput(3)).toEqual([
+      "'second foreground cell'",
+    ]);
+  });
 });