first comit

venv/lib/python3.10/site-packages/jupyter_client/ioloop/manager.py

@@ -0,0 +1,116 @@
+"""A kernel manager with a tornado IOLoop"""
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+import typing as t
+
+import zmq
+from tornado import ioloop
+from traitlets import Instance, Type
+from zmq.eventloop.zmqstream import ZMQStream
+
+from ..manager import AsyncKernelManager, KernelManager
+from .restarter import AsyncIOLoopKernelRestarter, IOLoopKernelRestarter
+
+
+def as_zmqstream(f: t.Any) -> t.Callable:
+    """Convert a socket to a zmq stream."""
+
+    def wrapped(self: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        save_socket_class = None
+        # zmqstreams only support sync sockets
+        if self.context._socket_class is not zmq.Socket:
+            save_socket_class = self.context._socket_class
+            self.context._socket_class = zmq.Socket
+        try:
+            socket = f(self, *args, **kwargs)
+        finally:
+            if save_socket_class:
+                # restore default socket class
+                self.context._socket_class = save_socket_class
+        return ZMQStream(socket, self.loop)
+
+    return wrapped
+
+
+class IOLoopKernelManager(KernelManager):
+    """An io loop kernel manager."""
+
+    loop = Instance("tornado.ioloop.IOLoop")
+
+    def _loop_default(self) -> ioloop.IOLoop:
+        return ioloop.IOLoop.current()
+
+    restarter_class = Type(
+        default_value=IOLoopKernelRestarter,
+        klass=IOLoopKernelRestarter,
+        help=(
+            "Type of KernelRestarter to use. "
+            "Must be a subclass of IOLoopKernelRestarter.\n"
+            "Override this to customize how kernel restarts are managed."
+        ),
+        config=True,
+    )
+    _restarter: t.Any = Instance("jupyter_client.ioloop.IOLoopKernelRestarter", allow_none=True)
+
+    def start_restarter(self) -> None:
+        """Start the restarter."""
+        if self.autorestart and self.has_kernel:
+            if self._restarter is None:
+                self._restarter = self.restarter_class(
+                    kernel_manager=self, loop=self.loop, parent=self, log=self.log
+                )
+            self._restarter.start()
+
+    def stop_restarter(self) -> None:
+        """Stop the restarter."""
+        if self.autorestart and self._restarter is not None:
+            self._restarter.stop()
+
+    connect_shell = as_zmqstream(KernelManager.connect_shell)
+    connect_control = as_zmqstream(KernelManager.connect_control)
+    connect_iopub = as_zmqstream(KernelManager.connect_iopub)
+    connect_stdin = as_zmqstream(KernelManager.connect_stdin)
+    connect_hb = as_zmqstream(KernelManager.connect_hb)
+
+
+class AsyncIOLoopKernelManager(AsyncKernelManager):
+    """An async ioloop kernel manager."""
+
+    loop = Instance("tornado.ioloop.IOLoop")
+
+    def _loop_default(self) -> ioloop.IOLoop:
+        return ioloop.IOLoop.current()
+
+    restarter_class = Type(
+        default_value=AsyncIOLoopKernelRestarter,
+        klass=AsyncIOLoopKernelRestarter,
+        help=(
+            "Type of KernelRestarter to use. "
+            "Must be a subclass of AsyncIOLoopKernelManager.\n"
+            "Override this to customize how kernel restarts are managed."
+        ),
+        config=True,
+    )
+    _restarter: t.Any = Instance(
+        "jupyter_client.ioloop.AsyncIOLoopKernelRestarter", allow_none=True
+    )
+
+    def start_restarter(self) -> None:
+        """Start the restarter."""
+        if self.autorestart and self.has_kernel:
+            if self._restarter is None:
+                self._restarter = self.restarter_class(
+                    kernel_manager=self, loop=self.loop, parent=self, log=self.log
+                )
+            self._restarter.start()
+
+    def stop_restarter(self) -> None:
+        """Stop the restarter."""
+        if self.autorestart and self._restarter is not None:
+            self._restarter.stop()
+
+    connect_shell = as_zmqstream(AsyncKernelManager.connect_shell)
+    connect_control = as_zmqstream(AsyncKernelManager.connect_control)
+    connect_iopub = as_zmqstream(AsyncKernelManager.connect_iopub)
+    connect_stdin = as_zmqstream(AsyncKernelManager.connect_stdin)
+    connect_hb = as_zmqstream(AsyncKernelManager.connect_hb)