optionally choose new random ports on kernel restart

Author: rolweber

jupyter_client/connect.py

@@ -423,7 +423,8 @@ def cleanup_ipc_files(self):
     def _record_random_port_names(self):
         """Records which of the ports are randomly assigned.
 
-        Records on first invocation. Does nothing on later invocations."""
+        Records on first invocation, if the transport is tcp.
+        Does nothing on later invocations."""
 
         if self.transport != 'tcp':
             return
@@ -435,6 +436,20 @@ def _record_random_port_names(self):
             if getattr(self, name) <= 0:
                 self._random_port_names.append(name)
 
+    def cleanup_random_ports(self):
+        """Forgets randomly assigned port numbers and cleans up the connection file.
+
+        Does nothing if no port numbers have been randomly assigned.
+        In particular, does nothing unless the transport is tcp.
+        """
+
+        if not self._random_port_names:
+            return
+
+        for name in self._random_port_names:
+            setattr(self, name, 0)
+
+        self.cleanup_connection_file()
 
     def write_connection_file(self):
         """Write connection info to JSON dict in self.connection_file."""

jupyter_client/manager.py

@@ -326,12 +326,9 @@ def shutdown_kernel(self, now=False, restart=False):
 
         self.cleanup(connection_file=not restart)
 
-    def restart_kernel(self, now=False, **kw):
+    def restart_kernel(self, now=False, newports=False, **kw):
         """Restarts a kernel with the arguments that were used to launch it.
 
-        If the old kernel was launched with random ports, the same ports will be
-        used for the new kernel. The same connection file is used again.
-
         Parameters
         ----------
         now : bool, optional
@@ -342,6 +339,14 @@ def restart_kernel(self, now=False, **kw):
             In all cases the kernel is restarted, the only difference is whether
             it is given a chance to perform a clean shutdown or not.
 
+        newports : bool, optional
+            If the old kernel was launched with random ports, this flag decides
+            whether the same ports and connection file will be used again.
+            If False, the same ports and connection file are used. This is
+            the default. If True, new random port numbers are chosen and a
+            new connection file is written. It is still possible that the newly
+            chosen random port numbers happen to be the same as the old ones.
+
         `**kw` : optional
             Any options specified here will overwrite those used to launch the
             kernel.
@@ -353,6 +358,9 @@ def restart_kernel(self, now=False, **kw):
             # Stop currently running kernel.
             self.shutdown_kernel(now=now, restart=True)
 
+            if newports:
+                self.cleanup_random_ports()
+
             # Start new kernel.
             self._launch_args.update(kw)
             self.start_kernel(**self._launch_args)

jupyter_client/tests/test_connect.py

@@ -183,6 +183,19 @@ def test_mixin_record_random_ports():
     with TemporaryDirectory() as d:
         dc = DummyConfigurable(data_dir=d, kernel_name='via-tcp', transport='tcp')
         dc.write_connection_file()
+
         assert dc._connection_file_written
         assert os.path.exists(dc.connection_file)
         assert dc._random_port_names == connect.port_names
+
+
+def test_mixin_cleanup_random_ports():
+    with TemporaryDirectory() as d:
+        dc = DummyConfigurable(data_dir=d, kernel_name='via-tcp', transport='tcp')
+        dc.write_connection_file()
+        filename = dc.connection_file
+        dc.cleanup_random_ports()
+
+        assert not os.path.exists(filename)
+        for name in dc._random_port_names:
+            assert getattr(dc, name) == 0