Repository: jupyter/jupyter_client
Pull request #880 (2022-11-15 00:35:41)

[pre-commit.ci] pre-commit autoupdate
<!--pre-commit.ci start-->
updates:
- [github.com/asottile/pyupgrade: v3.2.0 → v3.2.2](https://github.com/asottile/pyupgrade/compare/v3.2.0...v3.2.2)
- [github.com/pre-commit/mirrors-mypy: v0.982 → v0.990](https://github.com/pre-commit/mirrors-mypy/compare/v0.982...v0.990)
- [github.com/sirosen/check-jsonschema: 0.18.4 → 0.19.1](https://github.com/sirosen/check-jsonschema/compare/0.18.4...0.19.1)
<!--pre-commit.ci end-->


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ca7c0d5..40a147e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -38,7 +38,7 @@ repos:
       - id: mdformat
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.2.0
+    rev: v3.2.2
     hooks:
       - id: pyupgrade
         args: [--py38-plus]
@@ -60,7 +60,7 @@ repos:
         stages: [manual]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.982
+    rev: v0.990
     hooks:
       - id: mypy
         exclude: tests
@@ -75,7 +75,7 @@ repos:
         args: [--max-line-length=100]
 
   - repo: https://github.com/sirosen/check-jsonschema
-    rev: 0.18.4
+    rev: 0.19.1
     hooks:
       - id: check-jsonschema
         name: "Check GitHub Workflows"
diff --git a/jupyter_client/connect.py b/jupyter_client/connect.py
index e23ef97..97bab17 100644
--- a/jupyter_client/connect.py
+++ b/jupyter_client/connect.py
@@ -157,21 +157,9 @@ def write_connection_file(
     cfg["signature_scheme"] = signature_scheme
     cfg["kernel_name"] = kernel_name
 
-    # Prevent over-writing a file that has already been written with the same
-    # info.  This is to prevent a race condition where the process has
-    # already been launched but has not yet read the connection file.
-    if os.path.exists(fname):
-        with open(fname) as f:
-            try:
-                data = json.load(f)
-                if data == cfg:
-                    return fname, cfg
-            except Exception:
-                pass
-
     # Only ever write this file as user read/writeable
     # This would otherwise introduce a vulnerability as a file has secrets
-    # which would let others execute arbitrarily code as you
+    # which would let others execute arbitrary code as you
     with secure_write(fname) as f:
         f.write(json.dumps(cfg, indent=2))
 
@@ -580,18 +568,70 @@ class ConnectionFileMixin(LoggingConfigurable):
         if "signature_scheme" in info:
             self.session.signature_scheme = info["signature_scheme"]
 
-    def _force_connection_info(self, info: KernelConnectionInfo) -> None:
-        """Unconditionally loads connection info from a dict containing connection info.
+    def _reconcile_connection_info(self, info: KernelConnectionInfo) -> None:
+        """Reconciles the connection information returned from the Provisioner.
 
-        Overwrites connection info-based attributes, regardless of their current values
-        and writes this information to the connection file.
+        Because some provisioners (like derivations of LocalProvisioner) may have already
+        written the connection file, this method needs to ensure that, if the connection
+        file exists, its contents match that of what was returned by the provisioner.  If
+        the file does exist and its contents do not match, a ValueError is raised.
+
+        If the file does not exist, the connection information in 'info' is loaded into the
+        KernelManager and written to the file.
         """
-        # Reset current ports to 0 and indicate file has not been written to enable override
-        self._connection_file_written = False
-        for name in port_names:
-            setattr(self, name, 0)
-        self.load_connection_info(info)
-        self.write_connection_file()
+        # Prevent over-writing a file that has already been written with the same
+        # info.  This is to prevent a race condition where the process has
+        # already been launched but has not yet read the connection file - as is
+        # the case with LocalProvisioners.
+        file_exists: bool = False
+        if os.path.exists(self.connection_file):
+            with open(self.connection_file) as f:
+                file_info = json.load(f)
+                # Prior to the following comparison, we need to adjust the value of "key" to
+                # be bytes, otherwise the comparison below will fail.
+                file_info["key"] = file_info["key"].encode()
+                if not self._equal_connections(info, file_info):
+                    raise ValueError(
+                        "Connection file already exists and does not match "
+                        "the expected values returned from provisioner!"
+                    )
+                file_exists = True
+
+        if not file_exists:
+            # Load the connection info and write out file. Note, this does not necessarily
+            # overwrite non-zero port values, so we'll validate afterward.
+            self.load_connection_info(info)
+            self.write_connection_file()
+
+        # Ensure what is in KernelManager is what we expect.  This will catch issues if the file
+        # already existed, yet it's contents differed from the KernelManager's (and provisioner).
+        km_info = self.get_connection_info()
+        if not self._equal_connections(info, km_info):
+            raise ValueError(
+                "KernelManager's connection information already exists and does not match "
+                "the expected values returned from provisioner!"
+            )
+
+    @staticmethod
+    def _equal_connections(conn1: KernelConnectionInfo, conn2: KernelConnectionInfo) -> bool:
+        """Compares pertinent keys of connection info data. Returns True if equivalent, False otherwise."""
+
+        pertinent_keys = [
+            "key",
+            "ip",
+            "stdin_port",
+            "iopub_port",
+            "shell_port",
+            "control_port",
+            "hb_port",
+            "transport",
+            "signature_scheme",
+        ]
+
+        for key in pertinent_keys:
+            if conn1.get(key) != conn2.get(key):
+                return False
+        return True
 
     # --------------------------------------------------------------------------
     # Creating connected sockets
diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py
index f6257f5..e572566 100644
--- a/jupyter_client/manager.py
+++ b/jupyter_client/manager.py
@@ -309,8 +309,9 @@ class KernelManager(ConnectionFileMixin):
         assert self.provisioner is not None
         connection_info = await self.provisioner.launch_kernel(kernel_cmd, **kw)
         assert self.provisioner.has_process
-        # Provisioner provides the connection information.  Load into kernel manager and write file.
-        self._force_connection_info(connection_info)
+        # Provisioner provides the connection information.  Load into kernel manager
+        # and write the connection file, if not already done.
+        self._reconcile_connection_info(connection_info)
 
     _launch_kernel = run_sync(_async_launch_kernel)
 
diff --git a/jupyter_client/multikernelmanager.py b/jupyter_client/multikernelmanager.py
index 8048e47..f5352f9 100644
--- a/jupyter_client/multikernelmanager.py
+++ b/jupyter_client/multikernelmanager.py
@@ -381,7 +381,7 @@ class MultiKernelManager(LoggingConfigurable):
     restart_kernel = run_sync(_async_restart_kernel)
 
     @kernel_method
-    def is_alive(self, kernel_id: str) -> bool:
+    def is_alive(self, kernel_id: str) -> bool:  # type:ignore[empty-body]
         """Is the kernel alive.
 
         This calls KernelManager.is_alive() which calls Popen.poll on the
@@ -422,7 +422,7 @@ class MultiKernelManager(LoggingConfigurable):
         """remove a callback for the KernelRestarter"""
 
     @kernel_method
-    def get_connection_info(self, kernel_id: str) -> t.Dict[str, t.Any]:
+    def get_connection_info(self, kernel_id: str) -> t.Dict[str, t.Any]:  # type:ignore[empty-body]
         """Return a dictionary of connection data for a kernel.
 
         Parameters
@@ -440,7 +440,9 @@ class MultiKernelManager(LoggingConfigurable):
         """
 
     @kernel_method
-    def connect_iopub(self, kernel_id: str, identity: t.Optional[bytes] = None) -> socket.socket:
+    def connect_iopub(  # type:ignore[empty-body]
+        self, kernel_id: str, identity: t.Optional[bytes] = None
+    ) -> socket.socket:
         """Return a zmq Socket connected to the iopub channel.
 
         Parameters
@@ -456,7 +458,9 @@ class MultiKernelManager(LoggingConfigurable):
         """
 
     @kernel_method
-    def connect_shell(self, kernel_id: str, identity: t.Optional[bytes] = None) -> socket.socket:
+    def connect_shell(  # type:ignore[empty-body]
+        self, kernel_id: str, identity: t.Optional[bytes] = None
+    ) -> socket.socket:
         """Return a zmq Socket connected to the shell channel.
 
         Parameters
@@ -472,7 +476,9 @@ class MultiKernelManager(LoggingConfigurable):
         """
 
     @kernel_method
-    def connect_control(self, kernel_id: str, identity: t.Optional[bytes] = None) -> socket.socket:
+    def connect_control(  # type:ignore[empty-body]
+        self, kernel_id: str, identity: t.Optional[bytes] = None
+    ) -> socket.socket:
         """Return a zmq Socket connected to the control channel.
 
         Parameters
@@ -488,7 +494,9 @@ class MultiKernelManager(LoggingConfigurable):
         """
 
     @kernel_method
-    def connect_stdin(self, kernel_id: str, identity: t.Optional[bytes] = None) -> socket.socket:
+    def connect_stdin(  # type:ignore[empty-body]
+        self, kernel_id: str, identity: t.Optional[bytes] = None
+    ) -> socket.socket:
         """Return a zmq Socket connected to the stdin channel.
 
         Parameters
@@ -504,7 +512,9 @@ class MultiKernelManager(LoggingConfigurable):
         """
 
     @kernel_method
-    def connect_hb(self, kernel_id: str, identity: t.Optional[bytes] = None) -> socket.socket:
+    def connect_hb(  # type:ignore[empty-body]
+        self, kernel_id: str, identity: t.Optional[bytes] = None
+    ) -> socket.socket:
         """Return a zmq Socket connected to the hb channel.
 
         Parameters
diff --git a/jupyter_client/provisioning/local_provisioner.py b/jupyter_client/provisioning/local_provisioner.py
index 875995f..8d66f03 100644
--- a/jupyter_client/provisioning/local_provisioner.py
+++ b/jupyter_client/provisioning/local_provisioner.py
@@ -18,7 +18,7 @@ from ..localinterfaces import local_ips
 from .provisioner_base import KernelProvisionerBase
 
 
-class LocalProvisioner(KernelProvisionerBase):
+class LocalProvisioner(KernelProvisionerBase):  # type:ignore[misc]
     """
     :class:`LocalProvisioner` is a concrete class of ABC :py:class:`KernelProvisionerBase`
     and is the out-of-box default implementation used when no kernel provisioner is
diff --git a/jupyter_client/session.py b/jupyter_client/session.py
index d499a7a..c6f7ab4 100644
--- a/jupyter_client/session.py
+++ b/jupyter_client/session.py
@@ -834,7 +834,7 @@ class Session(Configurable):
             else:
                 try:
                     # check to see if buf supports the buffer protocol.
-                    view = memoryview(buf)
+                    view = memoryview(buf)  # type:ignore[assignment]
                 except TypeError as e:
                     raise TypeError("Buffer objects must support the buffer protocol.") from e
             # memoryview.contiguous is new in 3.3,


Test patch:
diff --git a/tests/test_connect.py b/tests/test_connect.py
index 8a237de..7103cf2 100644
--- a/tests/test_connect.py
+++ b/tests/test_connect.py
@@ -5,11 +5,13 @@ import json
 import os
 from tempfile import TemporaryDirectory
 
+import pytest
 from jupyter_core.application import JupyterApp
 from jupyter_core.paths import jupyter_runtime_dir
 
 from jupyter_client import connect
 from jupyter_client import KernelClient
+from jupyter_client import KernelManager
 from jupyter_client.consoleapp import JupyterConsoleApp
 from jupyter_client.session import Session
 
@@ -235,3 +237,67 @@ def test_mixin_cleanup_random_ports():
         assert not os.path.exists(filename)
         for name in dc._random_port_names:
             assert getattr(dc, name) == 0
+
+
+param_values = [
+    (True, True, None),
+    (True, False, ValueError),
+    (False, True, None),
+    (False, False, ValueError),
+]
+
+
+@pytest.mark.parametrize("file_exists, km_matches, expected_exception", param_values)
+def test_reconcile_connection_info(file_exists, km_matches, expected_exception):
+
+    expected_info = sample_info
+    mismatched_info = sample_info.copy()
+    mismatched_info["key"] = b"def456"
+    mismatched_info["shell_port"] = expected_info["shell_port"] + 42
+    mismatched_info["control_port"] = expected_info["control_port"] + 42
+
+    with TemporaryDirectory() as connection_dir:
+
+        cf = os.path.join(connection_dir, "kernel.json")
+        km = KernelManager()
+        km.connection_file = cf
+
+        if file_exists:
+            _, info = connect.write_connection_file(cf, **expected_info)
+            info["key"] = info["key"].encode()  # set 'key' back to bytes
+
+            if km_matches:
+                # Let this be the case where the connection file exists, and the KM has matching
+                # values prior to reconciliation.  This is the LocalProvisioner case.
+                provisioner_info = info
+                km.load_connection_info(provisioner_info)
+            else:
+                # Let this be the case where the connection file exists, the KM has no values
+                # prior to reconciliation, but the provisioner has returned different values
+                # and a ValueError is expected.
+                provisioner_info = mismatched_info
+        else:  # connection file does not exist
+            if km_matches:
+                # Let this be the case where the connection file does not exist, NOR does the KM
+                # have any values of its own and reconciliation sets those values. This is the
+                # non-LocalProvisioner case.
+                provisioner_info = expected_info
+            else:
+                # Let this be the case where the connection file does not exist, yet the KM
+                # has values that do not match those returned from the provisioner and a
+                # ValueError is expected.
+                km.load_connection_info(expected_info)
+                provisioner_info = mismatched_info
+
+        if expected_exception is None:
+            km._reconcile_connection_info(provisioner_info)
+            km_info = km.get_connection_info()
+            assert km._equal_connections(km_info, provisioner_info)
+        else:
+            with pytest.raises(expected_exception) as ex:
+                km._reconcile_connection_info(provisioner_info)
+            if file_exists:
+                assert "Connection file already exists" in str(ex.value)
+            else:
+                assert "KernelManager's connection information already exists" in str(ex.value)
+            assert km._equal_connections(km.get_connection_info(), provisioner_info) is False