Extend argument handling of do_execute with cell metadata (#1169)

Author: jjvraw

ipykernel/ipkernel.py

@@ -23,7 +23,7 @@
 from .debugger import Debugger, _is_debugpy_available
 from .eventloops import _use_appnope
 from .kernelbase import Kernel as KernelBase
-from .kernelbase import _accepts_cell_id
+from .kernelbase import _accepts_parameters
 from .zmqshell import ZMQInteractiveShell
 
 try:
@@ -347,6 +347,7 @@ async def do_execute(
         user_expressions=None,
         allow_stdin=False,
         *,
+        cell_meta=None,
         cell_id=None,
     ):
         """Handle code execution."""
@@ -359,7 +360,7 @@ async def do_execute(
         if hasattr(shell, "run_cell_async") and hasattr(shell, "should_run_async"):
             run_cell = shell.run_cell_async
             should_run_async = shell.should_run_async
-            with_cell_id = _accepts_cell_id(run_cell)
+            accepts_params = _accepts_parameters(run_cell, ["cell_id"])
         else:
             should_run_async = lambda cell: False  # noqa
             # older IPython,
@@ -368,7 +369,7 @@ async def do_execute(
             async def run_cell(*args, **kwargs):
                 return shell.run_cell(*args, **kwargs)
 
-            with_cell_id = _accepts_cell_id(shell.run_cell)
+            accepts_params = _accepts_parameters(shell.run_cell, ["cell_id"])
         try:
             # default case: runner is asyncio and asyncio is already running
             # TODO: this should check every case for "are we inside the runner",
@@ -390,7 +391,7 @@ async def run_cell(*args, **kwargs):
                     preprocessing_exc_tuple=preprocessing_exc_tuple,
                 )
             ):
-                if with_cell_id:
+                if accepts_params["cell_id"]:
                     coro = run_cell(
                         code,
                         store_history=store_history,
@@ -422,7 +423,7 @@ async def run_cell(*args, **kwargs):
                 # runner isn't already running,
                 # make synchronous call,
                 # letting shell dispatch to loop runners
-                if with_cell_id:
+                if accepts_params["cell_id"]:
                     res = shell.run_cell(
                         code,
                         store_history=store_history,

ipykernel/kernelbase.py

@@ -63,12 +63,18 @@
 from ._version import kernel_protocol_version
 
 
-def _accepts_cell_id(meth):
+def _accepts_parameters(meth, param_names):
     parameters = inspect.signature(meth).parameters
-    cid_param = parameters.get("cell_id")
-    return (cid_param and cid_param.kind == cid_param.KEYWORD_ONLY) or any(
-        p.kind == p.VAR_KEYWORD for p in parameters.values()
-    )
+    accepts = {param: False for param in param_names}
+
+    for param in param_names:
+        param_spec = parameters.get(param)
+        accepts[param] = (
+            param_spec
+            and param_spec.kind in [param_spec.KEYWORD_ONLY, param_spec.POSITIONAL_OR_KEYWORD]
+        ) or any(p.kind == p.VAR_KEYWORD for p in parameters.values())
+
+    return accepts
 
 
 class Kernel(SingletonConfigurable):
@@ -735,25 +741,28 @@ async def execute_request(self, stream, ident, parent):
             self.execution_count += 1
             self._publish_execute_input(code, parent, self.execution_count)
 
-        cell_id = (parent.get("metadata") or {}).get("cellId")
+        cell_meta = parent.get("metadata", {})
+        cell_id = cell_meta.get("cellId")
 
-        if _accepts_cell_id(self.do_execute):
-            reply_content = self.do_execute(
-                code,
-                silent,
-                store_history,
-                user_expressions,
-                allow_stdin,
-                cell_id=cell_id,
-            )
-        else:
-            reply_content = self.do_execute(
-                code,
-                silent,
-                store_history,
-                user_expressions,
-                allow_stdin,
-            )
+        # Check which parameters do_execute can accept
+        accepts_params = _accepts_parameters(self.do_execute, ["cell_meta", "cell_id"])
+
+        # Arguments based on the do_execute signature
+        do_execute_args = {
+            "code": code,
+            "silent": silent,
+            "store_history": store_history,
+            "user_expressions": user_expressions,
+            "allow_stdin": allow_stdin,
+        }
+
+        if accepts_params["cell_meta"]:
+            do_execute_args["cell_meta"] = cell_meta
+        if accepts_params["cell_id"]:
+            do_execute_args["cell_id"] = cell_id
+
+        # Call do_execute with the appropriate arguments
+        reply_content = self.do_execute(**do_execute_args)
 
         if inspect.isawaitable(reply_content):
             reply_content = await reply_content
@@ -793,6 +802,7 @@ def do_execute(
         user_expressions=None,
         allow_stdin=False,
         *,
+        cell_meta=None,
         cell_id=None,
     ):
         """Execute user code. Must be overridden by subclasses."""