1  """Terminal input and output prompts."""
   2  
   3  from pygments.token import Token
   4  import sys
   5  
   6  from IPython.core.displayhook import DisplayHook
   7  
   8  from prompt_toolkit.formatted_text import fragment_list_width, PygmentsTokens
   9  from prompt_toolkit.shortcuts import print_formatted_text
  10  from prompt_toolkit.enums import EditingMode
  11  
  12  
  13  class Prompts(object):
  14      def __init__(self, shell):
  15          self.shell = shell
  16  
  17      def vi_mode(self):
  18          if (getattr(self.shell.pt_app, 'editing_mode', None) == EditingMode.VI
  19                  and self.shell.prompt_includes_vi_mode):
  20              mode = str(self.shell.pt_app.app.vi_state.input_mode)
  21              if mode.startswith('InputMode.'):
  22                  mode = mode[10:13].lower()
  23              elif mode.startswith('vi-'):
  24                  mode = mode[3:6]
  25              return '['+mode+'] '
  26          return ''
  27  
  28      def current_line(self) -> int:
  29          if self.shell.pt_app is not None:
  30              return self.shell.pt_app.default_buffer.document.cursor_position_row or 0
  31          return 0
  32  
  33      def in_prompt_tokens(self):
  34          return [
  35              (Token.Prompt, self.vi_mode()),
  36              (
  37                  Token.Prompt,
  38                  self.shell.prompt_line_number_format.format(
  39                      line=1, rel_line=-self.current_line()
  40                  ),
  41              ),
  42              (Token.Prompt, "In ["),
  43              (Token.PromptNum, str(self.shell.execution_count)),
  44              (Token.Prompt, ']: '),
  45          ]
  46  
  47      def _width(self):
  48          return fragment_list_width(self.in_prompt_tokens())
  49  
  50      def continuation_prompt_tokens(self, width=None, *, lineno=None):
  51          if width is None:
  52              width = self._width()
  53          line = lineno + 1 if lineno is not None else 0
  54          prefix = " " * len(
  55              self.vi_mode()
  56          ) + self.shell.prompt_line_number_format.format(
  57              line=line, rel_line=line - self.current_line() - 1
  58          )
  59          return [
  60              (
  61                  Token.Prompt,
  62                  prefix + (" " * (width - len(prefix) - 5)) + "...: ",
  63              ),
  64          ]
  65  
  66      def rewrite_prompt_tokens(self):
  67          width = self._width()
  68          return [
  69              (Token.Prompt, ('-' * (width - 2)) + '> '),
  70          ]
  71  
  72      def out_prompt_tokens(self):
  73          return [
  74              (Token.OutPrompt, 'Out['),
  75              (Token.OutPromptNum, str(self.shell.execution_count)),
  76              (Token.OutPrompt, ']: '),
  77          ]
  78  
  79  class ClassicPrompts(Prompts):
  80      def in_prompt_tokens(self):
  81          return [
  82              (Token.Prompt, '>>> '),
  83          ]
  84  
  85      def continuation_prompt_tokens(self, width=None):
  86          return [
  87              (Token.Prompt, '... ')
  88          ]
  89  
  90      def rewrite_prompt_tokens(self):
  91          return []
  92  
  93      def out_prompt_tokens(self):
  94          return []
  95  
  96  class RichPromptDisplayHook(DisplayHook):
  97      """Subclass of base display hook using coloured prompt"""
  98      def write_output_prompt(self):
  99          sys.stdout.write(self.shell.separate_out)
 100          # If we're not displaying a prompt, it effectively ends with a newline,
 101          # because the output will be left-aligned.
 102          self.prompt_end_newline = True
 103  
 104          if self.do_full_cache:
 105              tokens = self.shell.prompts.out_prompt_tokens()
 106              prompt_txt = "".join(s for _, s in tokens)
 107              if prompt_txt and not prompt_txt.endswith("\n"):
 108                  # Ask for a newline before multiline output
 109                  self.prompt_end_newline = False
 110  
 111              if self.shell.pt_app:
 112                  print_formatted_text(PygmentsTokens(tokens),
 113                      style=self.shell.pt_app.app.style, end='',
 114                  )
 115              else:
 116                  sys.stdout.write(prompt_txt)
 117  
 118      def write_format_data(self, format_dict, md_dict=None) -> None:
 119          assert self.shell is not None
 120          if self.shell.mime_renderers:
 121  
 122              for mime, handler in self.shell.mime_renderers.items():
 123                  if mime in format_dict:
 124                      handler(format_dict[mime], None)
 125                      return
 126                  
 127          super().write_format_data(format_dict, md_dict)
 128  
 129