Detect subclass overrides of cleanup methods

jupyter_client/manager.py

@@ -388,10 +388,23 @@ def shutdown_kernel(self, now=False, restart=False):
             # most 1s, checking every 0.1s.
             self.finish_shutdown()
 
-        from . import __version__
-        from distutils.version import LooseVersion
-
-        if LooseVersion(__version__) < LooseVersion('6.2'):
+        # In 6.1.5, a new method, cleanup_resources(), was introduced to address
+        # a leak issue (https://github.com/jupyter/jupyter_client/pull/548) and
+        # replaced the existing cleanup() method.  However, that method introduction
+        # breaks subclass implementations that override cleanup() since it would
+        # circumvent cleanup() functionality implemented in subclasses.
+        # By detecting if the current instance overrides cleanup(), we can determine
+        # if the deprecated path of calling cleanup() should be performed - which avoids
+        # unnecessary deprecation warnings in a majority of configurations in which
+        # subclassed KernelManager instances are not in use.
+        # Note: because subclasses may have already implemented cleanup_resources()
+        # but need to support older jupyter_clients, we should only take the deprecated
+        # path if cleanup() is overridden but cleanup_resources() is not.
+
+        overrides_cleanup = type(self).cleanup is not KernelManager.cleanup
+        overrides_cleanup_resources = type(self).cleanup_resources is not KernelManager.cleanup_resources
+
+        if overrides_cleanup and not overrides_cleanup_resources:
             self.cleanup(connection_file=not restart)
         else:
             self.cleanup_resources(restart=restart)
@@ -609,10 +622,11 @@ async def shutdown_kernel(self, now=False, restart=False):
             # most 1s, checking every 0.1s.
             await self.finish_shutdown()
 
-        from . import __version__
-        from distutils.version import LooseVersion
+        # See comment in KernelManager.shutdown_kernel().
+        overrides_cleanup = type(self).cleanup is not AsyncKernelManager.cleanup
+        overrides_cleanup_resources = type(self).cleanup_resources is not AsyncKernelManager.cleanup_resources
 
-        if LooseVersion(__version__) < LooseVersion('6.2'):
+        if overrides_cleanup and not overrides_cleanup_resources:
             self.cleanup(connection_file=not restart)
         else:
             self.cleanup_resources(restart=restart)

jupyter_client/multikernelmanager.py

@@ -17,7 +17,7 @@
 from ipython_genutils.py3compat import unicode_type
 
 from .kernelspec import NATIVE_KERNEL_NAME, KernelSpecManager
-from .manager import AsyncKernelManager
+from .manager import KernelManager, AsyncKernelManager
 
 
 class DuplicateKernelError(Exception):
@@ -250,7 +250,18 @@ def shutdown_all(self, now=False):
             self.request_shutdown(kid)
         for kid in kids:
             self.finish_shutdown(kid)
-            self.cleanup(kid)
+
+            # Determine which cleanup method to call
+            # See comment in KernelManager.shutdown_kernel().
+            km = self.get_kernel(kid)
+            overrides_cleanup = type(km).cleanup is not KernelManager.cleanup
+            overrides_cleanup_resources = type(km).cleanup_resources is not KernelManager.cleanup_resources
+
+            if overrides_cleanup and not overrides_cleanup_resources:
+                km.cleanup(connection_file=True)
+            else:
+                km.cleanup_resources(restart=False)
+
             self.remove_kernel(kid)
 
     @kernel_method

jupyter_client/tests/test_kernelmanager.py

@@ -16,7 +16,7 @@
 import multiprocessing as mp
 import pytest
 from unittest import TestCase
-from tornado.testing import AsyncTestCase, gen_test, gen
+from tornado.testing import AsyncTestCase, gen_test
 
 from traitlets.config.loader import Config
 from jupyter_core import paths
@@ -210,7 +210,7 @@ def test_no_cleanup_shared_context(self):
         import zmq
         ctx = zmq.Context()
         km = KernelManager(context=ctx)
-        self.assertEquals(km.context, ctx)
+        self.assertEqual(km.context, ctx)
         self.assertIsNotNone(km.context)
 
         km.cleanup_resources(restart=False)
@@ -461,3 +461,102 @@ async def test_start_new_async_kernel(self):
             await km.shutdown_kernel(now=True)
             kc.stop_channels()
             self.assertTrue(km.context.closed)
+
+
+class AsyncKernelManagerSubclass(AsyncKernelManager):
+    """Used to test deprecation "routes" that are determined by superclass' detection of methods.
+
+       This class represents a current subclass that overrides both cleanup() and cleanup_resources()
+       in order to be compatible with older jupyter_clients.  We should find that cleanup_resources()
+       is called on these instances vix TestAsyncKernelManagerSubclass.
+    """
+
+    def cleanup(self, connection_file=True):
+        super(AsyncKernelManagerSubclass, self).cleanup(connection_file=connection_file)
+        self.which_cleanup = 'cleanup'
+
+    def cleanup_resources(self, restart=False):
+        super(AsyncKernelManagerSubclass, self).cleanup_resources(restart=restart)
+        self.which_cleanup = 'cleanup_resources'
+
+
+class AsyncKernelManagerWithCleanup(AsyncKernelManager):
+    """Used to test deprecation "routes" that are determined by superclass' detection of methods.
+
+       This class represents the older subclass that overrides cleanup().  We should find that
+       cleanup() is called on these instances via TestAsyncKernelManagerWithCleanup.
+    """
+
+    def cleanup(self, connection_file=True):
+        super(AsyncKernelManagerWithCleanup, self).cleanup(connection_file=connection_file)
+        self.which_cleanup = 'cleanup'
+
+
+class TestAsyncKernelManagerSubclass(AsyncTestCase):
+    def setUp(self):
+        super(TestAsyncKernelManagerSubclass, self).setUp()
+        self.env_patch = test_env()
+        self.env_patch.start()
+
+    def tearDown(self):
+        super(TestAsyncKernelManagerSubclass, self).tearDown()
+        self.env_patch.stop()
+
+    def _install_test_kernel(self):
+        kernel_dir = pjoin(paths.jupyter_data_dir(), 'kernels', 'signaltest')
+        os.makedirs(kernel_dir)
+        with open(pjoin(kernel_dir, 'kernel.json'), 'w') as f:
+            f.write(json.dumps({
+                'argv': [sys.executable,
+                         '-m', 'jupyter_client.tests.signalkernel',
+                         '-f', '{connection_file}'],
+                'display_name': "Signal Test Kernel",
+            }))
+
+    def _get_tcp_km(self):
+        c = Config()
+        km = AsyncKernelManagerSubclass(config=c)
+        return km
+
+    def _get_ipc_km(self):
+        c = Config()
+        c.KernelManager.transport = 'ipc'
+        c.KernelManager.ip = 'test'
+        km = AsyncKernelManagerSubclass(config=c)
+        return km
+
+    async def _run_lifecycle(self, km):
+        await km.start_kernel(stdout=PIPE, stderr=PIPE)
+        self.assertTrue(await km.is_alive())
+        await km.restart_kernel(now=True)
+        self.assertTrue(await km.is_alive())
+        await km.interrupt_kernel()
+        self.assertTrue(isinstance(km, AsyncKernelManager))
+        await km.shutdown_kernel(now=True)
+        self.assertFalse(await km.is_alive())
+        self.assertTrue(km.context.closed)
+
+    @gen_test
+    async def test_which_cleanup_method(self):
+        # This test confirms that the normal test operations run using cleanup_resources.
+        km = self._get_tcp_km()
+        await self._run_lifecycle(km)
+
+        if isinstance(km, AsyncKernelManagerSubclass):
+            assert km.which_cleanup == "cleanup_resources"
+        else:
+            assert km.which_cleanup == "cleanup"
+
+
+class TestAsyncKernelManagerWithCleanup(TestAsyncKernelManagerSubclass):
+    def _get_tcp_km(self):
+        c = Config()
+        km = AsyncKernelManagerWithCleanup(config=c)
+        return km
+
+    def _get_ipc_km(self):
+        c = Config()
+        c.KernelManager.transport = 'ipc'
+        c.KernelManager.ip = 'test'
+        km = AsyncKernelManagerWithCleanup(config=c)
+        return km