1  """ Implements a fully blocking kernel client.
   2  
   3  Useful for test suites and blocking terminal interfaces.
   4  """
   5  import sys
   6  
   7  # -----------------------------------------------------------------------------
   8  #  Copyright (C) 2012  The IPython Development Team
   9  #
  10  #  Distributed under the terms of the BSD License.  The full license is in
  11  #  the file LICENSE, distributed as part of this software.
  12  # -----------------------------------------------------------------------------
  13  from queue import Empty, Queue
  14  
  15  # IPython imports
  16  from traitlets import Type
  17  
  18  # Local imports
  19  from .channels import InProcessChannel
  20  from .client import InProcessKernelClient
  21  
  22  
  23  class BlockingInProcessChannel(InProcessChannel):
  24      """A blocking in-process channel."""
  25  
  26      def __init__(self, *args, **kwds):
  27          """Initialize the channel."""
  28          super().__init__(*args, **kwds)
  29          self._in_queue: Queue[object] = Queue()
  30  
  31      def call_handlers(self, msg):
  32          """Call the handlers for a message."""
  33          self._in_queue.put(msg)
  34  
  35      def get_msg(self, block=True, timeout=None):
  36          """Gets a message if there is one that is ready."""
  37          if timeout is None:
  38              # Queue.get(timeout=None) has stupid uninteruptible
  39              # behavior, so wait for a week instead
  40              timeout = 604800
  41          return self._in_queue.get(block, timeout)
  42  
  43      def get_msgs(self):
  44          """Get all messages that are currently ready."""
  45          msgs = []
  46          while True:
  47              try:
  48                  msgs.append(self.get_msg(block=False))
  49              except Empty:
  50                  break
  51          return msgs
  52  
  53      def msg_ready(self):
  54          """Is there a message that has been received?"""
  55          return not self._in_queue.empty()
  56  
  57  
  58  class BlockingInProcessStdInChannel(BlockingInProcessChannel):
  59      """A blocking in-process stdin channel."""
  60  
  61      def call_handlers(self, msg):
  62          """Overridden for the in-process channel.
  63  
  64          This methods simply calls raw_input directly.
  65          """
  66          msg_type = msg["header"]["msg_type"]
  67          if msg_type == "input_request":
  68              _raw_input = self.client.kernel._sys_raw_input
  69              prompt = msg["content"]["prompt"]
  70              print(prompt, end="", file=sys.__stdout__)
  71              sys.__stdout__.flush()
  72              self.client.input(_raw_input())
  73  
  74  
  75  class BlockingInProcessKernelClient(InProcessKernelClient):
  76      """A blocking in-process kernel client."""
  77  
  78      # The classes to use for the various channels.
  79      shell_channel_class = Type(BlockingInProcessChannel)  # type:ignore[arg-type]
  80      iopub_channel_class = Type(BlockingInProcessChannel)  # type:ignore[arg-type]
  81      stdin_channel_class = Type(BlockingInProcessStdInChannel)  # type:ignore[arg-type]
  82  
  83      def wait_for_ready(self):
  84          """Wait for kernel info reply on shell channel."""
  85          while True:
  86              self.kernel_info()
  87              try:
  88                  msg = self.shell_channel.get_msg(block=True, timeout=1)
  89              except Empty:
  90                  pass
  91              else:
  92                  if msg["msg_type"] == "kernel_info_reply":
  93                      # Checking that IOPub is connected. If it is not connected, start over.
  94                      try:
  95                          self.iopub_channel.get_msg(block=True, timeout=0.2)
  96                      except Empty:
  97                          pass
  98                      else:
  99                          self._handle_kernel_info_reply(msg)
 100                          break
 101  
 102          # Flush IOPub channel
 103          while True:
 104              try:
 105                  msg = self.iopub_channel.get_msg(block=True, timeout=0.2)
 106                  print(msg["msg_type"])
 107              except Empty:
 108                  break
 109