From 49dbf0f83d7e146c7036c054a8564b30a510c49c Mon Sep 17 00:00:00 2001
From: Matthias Bussonnier <bussonniermatthias@gmail.com>
Date: Mon, 23 Oct 2023 18:18:40 +0200
Subject: [PATCH] Allow to update kernels env in between restart. (#987)

This allow to update the running env of kernels.

This will allow resolves some inconsistencies between kernel restart, and actually
stopping and starting a kernel session.

In parsley via a few lines in Jupyter Server in update_session, this
will let us update the `__session__` which is supposed to reflect some
information about the current session, and help correct
ipython/ipykernel#1102 where renaming a notebook is not reflected in the
kernel unless you manually stop/start the session, or restart the
server.
---
 jupyter_client/manager.py            | 15 +++++++++++++++
 jupyter_client/multikernelmanager.py | 11 +++++++++++
 2 files changed, 26 insertions(+)

diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py
index 3dff5433..f046f5a6 100644
--- a/jupyter_client/manager.py
+++ b/jupyter_client/manager.py
@@ -268,6 +268,21 @@ def client(self, **kwargs: t.Any) -> BlockingKernelClient:
     # Kernel management
     # --------------------------------------------------------------------------
 
+    def update_env(self, *, env: t.Dict[str, str]) -> None:
+        """
+        Allow to update the environment of a kernel manager.
+
+        This will take effect only after kernel restart when the new env is
+        passed to the new kernel.
+
+        This is useful as some of the information of the current kernel reflect
+        the state of the session that started it, and those session information
+        (like the attach file path, or name), are mutable.
+
+        .. version-added: 8.5
+        """
+        self._launch_args['env'].update(env)
+
     def format_kernel_cmd(self, extra_arguments: t.Optional[t.List[str]] = None) -> t.List[str]:
         """Replace templated args (e.g. {connection_file})"""
         extra_arguments = extra_arguments or []
diff --git a/jupyter_client/multikernelmanager.py b/jupyter_client/multikernelmanager.py
index 2ebd0e9d..a3799b2d 100644
--- a/jupyter_client/multikernelmanager.py
+++ b/jupyter_client/multikernelmanager.py
@@ -212,6 +212,17 @@ def pre_start_kernel(
         )
         return km, kernel_name, kernel_id
 
+    def update_env(self, *, kernel_id: str, env: t.Dict[str, str]) -> None:
+        """
+        Allow to update the environment of the given kernel.
+
+        Forward the update env request to the corresponding kernel.
+
+        .. version-added: 8.5
+        """
+        if kernel_id in self:
+            self._kernels[kernel_id].update_env(env)
+
     async def _add_kernel_when_ready(
         self, kernel_id: str, km: KernelManager, kernel_awaitable: t.Awaitable
     ) -> None: