Repository: jupyter/jupyter_client
Pull request #989 (2023-10-25 09:10:08)

BUG: Fix Kwarg only in update_env
closes jupyterlab/jupyterlab#15301

blink1073: Looks like there's still an issue with the downstream server tests.

```python
   Traceback (most recent call last):
    File "/home/runner/test_venv/lib/python3.11/site-packages/tornado/web.py", line 1786, in _execute
      result = await result
               ^^^^^^^^^^^^
    File "/home/runner/test_venv/downstream_test/jupyter_server/jupyter_server/services/sessions/handlers.py", line 172, in patch
      await sm.update_session(session_id, **changes)
    File "/home/runner/test_venv/downstream_test/jupyter_server/jupyter_server/services/sessions/sessionmanager.py", line 469, in update_session
      self.kernel_manager.update_env(kernel_id=kernel_id, env=self.get_kernel_env(path, name))
    File "/home/runner/test_venv/lib/python3.11/site-packages/jupyter_client/multikernelmanager.py", line 226, in update_env
      self._kernels[kernel_id].update_env(env=env)
    File "/home/runner/test_venv/lib/python3.11/site-packages/jupyter_client/manager.py", line 284, in update_env
      self._launch_args['env'].update(env)
      ~~~~~~~~~~~~~~~~~^^^^^^^
  KeyError: 'env'
```
blink1073: I'm reverting this change in server for now.
Carreau: I pushed a check that env exists.
blink1073: Have you tested this with server to 2.8.1?
Carreau: Ok, I've been more agressive on the condition, checking that everything exists and are dicts, and added all the corresponding test cases.

I've tested on Jupyter_server  2.7.0, 2.9.1 (does not crash, but does nothing), 
and on 2.9.0 (it does work an update `__session__`). 
Carreau: Let me see if I can fix the typing of `_launch_args`, that is set to `Any`, I think we can make it t.Optional[Dict]
Carreau: ok, I can't make it ` t.Optional[Dict[str, Any]]`, as `Dict` is not subscriptable, but I did make it `t.Optional[Dict]`
blink1073: Papermill failure is known: https://github.com/nteract/papermill/issues/737

Patch:
diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py
index 6131838..2f7fa0e 100644
--- a/jupyter_client/manager.py
+++ b/jupyter_client/manager.py
@@ -19,6 +19,7 @@ from jupyter_core.utils import run_sync
 from traitlets import (
     Any,
     Bool,
+    Dict,
     DottedObjectName,
     Float,
     Instance,
@@ -206,7 +207,7 @@ class KernelManager(ConnectionFileMixin):
         return self.kernel_name in {"python", "python2", "python3"}
 
     # Protected traits
-    _launch_args: Any = Any()
+    _launch_args: t.Optional["Dict[str, Any]"] = Dict(allow_none=True)
     _control_socket: Any = Any()
 
     _restarter: Any = Any()
@@ -281,7 +282,13 @@ class KernelManager(ConnectionFileMixin):
 
         .. version-added: 8.5
         """
-        self._launch_args["env"].update(env)
+        # Mypy think this is unreachable as it see _launch_args as Dict, not t.Dict
+        if (
+            isinstance(self._launch_args, dict)
+            and "env" in self._launch_args
+            and isinstance(self._launch_args["env"], dict)  # type: ignore [unreachable]
+        ):
+            self._launch_args["env"].update(env)  # type: ignore [unreachable]
 
     def format_kernel_cmd(self, extra_arguments: t.Optional[t.List[str]] = None) -> t.List[str]:
         """Replace templated args (e.g. {connection_file})"""
@@ -307,13 +314,14 @@ class KernelManager(ConnectionFileMixin):
         # is not usable by non python kernels because the path is being rerouted when
         # inside of a store app.
         # See this bug here: https://bugs.python.org/issue41196
-        ns = {
+        ns: t.Dict[str, t.Any] = {
             "connection_file": os.path.realpath(self.connection_file),
             "prefix": sys.prefix,
         }
 
         if self.kernel_spec:  # type:ignore[truthy-bool]
             ns["resource_dir"] = self.kernel_spec.resource_dir
+        assert isinstance(self._launch_args, dict)
 
         ns.update(self._launch_args)
 
@@ -371,7 +379,8 @@ class KernelManager(ConnectionFileMixin):
         self.shutting_down = False
         self.kernel_id = self.kernel_id or kw.pop("kernel_id", str(uuid.uuid4()))
         # save kwargs for use in restart
-        self._launch_args = kw.copy()
+        # assigning Traitlets Dicts to Dict make mypy unhappy but is ok
+        self._launch_args = kw.copy()  # type:ignore [assignment]
         if self.provisioner is None:  # will not be None on restarts
             self.provisioner = KPF.instance(parent=self.parent).create_provisioner_instance(
                 self.kernel_id,
diff --git a/jupyter_client/multikernelmanager.py b/jupyter_client/multikernelmanager.py
index 122db74..d14a3f8 100644
--- a/jupyter_client/multikernelmanager.py
+++ b/jupyter_client/multikernelmanager.py
@@ -223,7 +223,7 @@ class MultiKernelManager(LoggingConfigurable):
         .. version-added: 8.5
         """
         if kernel_id in self:
-            self._kernels[kernel_id].update_env(env)
+            self._kernels[kernel_id].update_env(env=env)
 
     async def _add_kernel_when_ready(
         self, kernel_id: str, km: KernelManager, kernel_awaitable: t.Awaitable


Test patch:
diff --git a/tests/test_manager.py b/tests/test_manager.py
index e3d6ea2..23c246a 100644
--- a/tests/test_manager.py
+++ b/tests/test_manager.py
@@ -32,3 +32,33 @@ def test_connection_file_real_path():
         km._launch_args = {}
         cmds = km.format_kernel_cmd()
         assert cmds[4] == "foobar"
+
+
+def test_env_update_launch_args_not_set():
+    km = KernelManager()
+    km.update_env(env={"A": "A"})
+
+
+def test_env_update_launch_args_not_dict():
+    km = KernelManager()
+    km._launch_args = None
+    km.update_env(env={"B": "B"})
+
+
+def test_env_update_launch_args_no_env():
+    km = KernelManager()
+    km._launch_args = {}
+    km.update_env(env={"C": "C"})
+
+
+def test_env_update_launch_args_env_not_dict():
+    km = KernelManager()
+    km._launch_args = {"env": None}
+    km.update_env(env={"D": "D"})
+
+
+def test_env_update_launch_args_env_dic():
+    km = KernelManager()
+    km._launch_args = {"env": {}}
+    km.update_env(env={"E": "E"})
+    assert km._launch_args["env"]["E"] == "E"