BUG: Fix Kwarg only in update_env (#989)

* BUG: Fix Kwarg only in update_env closes jupyterlab/jupyterlab#15301 * add typing informations * use str annotations
jupyter · Nov 6, 2023 · b61accf · b61accf
1 parent 79fa122
commit b61accf
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 5 deletions.
diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py
@@ -19,6 +19,7 @@
 from traitlets import (
     Any,
     Bool,
+    Dict,
     DottedObjectName,
     Float,
     Instance,
@@ -206,7 +207,7 @@ def ipykernel(self) -> bool:
         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 @@ def update_env(self, *, env: t.Dict[str, str]) -> None:
 
         .. 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 @@ def format_kernel_cmd(self, extra_arguments: t.Optional[t.List[str]] = None) ->
         # 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 @@ async def _async_pre_start_kernel(
         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
@@ -223,7 +223,7 @@ def update_env(self, *, kernel_id: str, env: t.Dict[str, str]) -> None:
         .. 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

diff --git 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"