feat: allow kwargs when writting connection_file (#953)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Steven Silvester <[email protected]>

Author: 3 people

jupyter_client/connect.py

@@ -40,6 +40,7 @@ def write_connection_file(
     transport: str = "tcp",
     signature_scheme: str = "hmac-sha256",
     kernel_name: str = "",
+    **kwargs: Any,
 ) -> Tuple[str, KernelConnectionInfo]:
     """Generates a JSON config file, including the selection of random ports.
 
@@ -140,6 +141,7 @@ def write_connection_file(
     cfg["transport"] = transport
     cfg["signature_scheme"] = signature_scheme
     cfg["kernel_name"] = kernel_name
+    cfg.update(kwargs)
 
     # Only ever write this file as user read/writeable
     # This would otherwise introduce a vulnerability as a file has secrets
@@ -483,7 +485,7 @@ def cleanup_random_ports(self) -> None:
 
         self.cleanup_connection_file()
 
-    def write_connection_file(self) -> None:
+    def write_connection_file(self, **kwargs: Any) -> None:
         """Write connection info to JSON dict in self.connection_file."""
         if self._connection_file_written and os.path.exists(self.connection_file):
             return
@@ -500,6 +502,7 @@ def write_connection_file(self) -> None:
             control_port=self.control_port,
             signature_scheme=self.session.signature_scheme,
             kernel_name=self.kernel_name,
+            **kwargs,
         )
         # write_connection_file also sets default ports:
         self._record_random_port_names()

jupyter_client/provisioning/local_provisioner.py

@@ -186,8 +186,11 @@ async def pre_launch(self, **kwargs: Any) -> Dict[str, Any]:
                 km.hb_port = lpc.find_available_port(km.ip)
                 km.control_port = lpc.find_available_port(km.ip)
                 self.ports_cached = True
-
-            km.write_connection_file()
+            if 'env' in kwargs:
+                jupyter_session = kwargs['env'].get("JPY_SESSION_NAME", "")
+                km.write_connection_file(jupyter_session=jupyter_session)
+            else:
+                km.write_connection_file()
             self.connection_info = km.get_connection_info()
 
             kernel_cmd = km.format_kernel_cmd(