From ea266fd6d4f463c70e181cc8b3c967c87106455e Mon Sep 17 00:00:00 2001
From: Tyler Pirtle <rtp@google.com>
Date: Fri, 13 Dec 2024 22:36:41 +0000
Subject: [PATCH] Add support for displaying image/png data.

This change maps image/png display data from execution results as <img> tags with in-line data.
---
 kernel_gateway/notebook_http/handlers.py | 40 +++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/kernel_gateway/notebook_http/handlers.py b/kernel_gateway/notebook_http/handlers.py
index a87c0eb..0ddc826 100644
--- a/kernel_gateway/notebook_http/handlers.py
+++ b/kernel_gateway/notebook_http/handlers.py
@@ -61,6 +61,32 @@ def initialize(self, sources, response_sources, kernel_pool, kernel_name, kernel
         self.response_sources = response_sources
         self.kernel_language = kernel_language
 
+    def _accumulate_display(self, results):
+        """Accumulates result chunks for "display" messages and prepares them as
+
+        in-line <img> tags. See
+        https://ipython.org/ipython-doc/3/development/messaging.html#display-data
+        for details on the display protocol.
+
+        Parameters
+        ----------
+        results: list
+          A list of results containing display data.
+        """
+        out = []
+        for result in results:
+            if "image/png" in result:
+                out.append(
+                    '<img alt="%s" src="data:image/png;base64,%s" />'
+                    % (result.get("text/plain", ""), result["image/png"])
+                )
+                continue
+            if "text/html" in result:
+                out.append(result["text/html"])
+            if "text/plain" in result:
+                out.append(result["text/plain"])
+        return out
+
     def finish_future(self, future, result_accumulator):
         """Resolves the promise to respond to a HTTP request handled by a
         kernel in the pool.
@@ -84,6 +110,10 @@ def finish_future(self, future, result_accumulator):
         """
         if result_accumulator["error"]:
             future.set_exception(CodeExecutionError(result_accumulator["error"]))
+        elif len(result_accumulator["display"]) > 0:
+            future.set_result(
+                "\n".join(self._accumulate_display(result_accumulator["display"]))
+            )
         elif len(result_accumulator["stream"]) > 0:
             future.set_result("".join(result_accumulator["stream"]))
         elif result_accumulator["result"]:
@@ -123,6 +153,9 @@ def on_recv(self, result_accumulator, future, parent_header, msg):
             # Store the execute result
             elif msg["header"]["msg_type"] == "execute_result":
                 result_accumulator["result"] = msg["content"]["data"]
+            # Accumulate display data
+            elif msg['header']['msg_type'] == 'display_data':
+                result_accumulator['display'].append(msg['content']['data'])
             # Accumulate the stream messages
             elif msg["header"]["msg_type"] == "stream":
                 # Only take stream output if it is on stdout or if the kernel
@@ -162,7 +195,12 @@ def execute_code(self, kernel_client, kernel_id, source_code):
             If the kernel returns any error
         """
         future = Future()
-        result_accumulator = {"stream": [], "error": None, "result": None}
+        result_accumulator = {
+            "display": [],
+            "stream": [],
+            "error": None,
+            "result": None,
+        }
         parent_header = kernel_client.execute(source_code)
         on_recv_func = partial(self.on_recv, result_accumulator, future, parent_header)
         self.kernel_pool.on_recv(kernel_id, on_recv_func)