prompts.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. """Terminal input and output prompts."""
  2. from __future__ import print_function
  3. from pygments.token import Token
  4. import sys
  5. from IPython.core.displayhook import DisplayHook
  6. from prompt_toolkit.layout.utils import token_list_width
  7. class Prompts(object):
  8. def __init__(self, shell):
  9. self.shell = shell
  10. def in_prompt_tokens(self, cli=None):
  11. return [
  12. (Token.Prompt, 'In ['),
  13. (Token.PromptNum, str(self.shell.execution_count)),
  14. (Token.Prompt, ']: '),
  15. ]
  16. def _width(self):
  17. return token_list_width(self.in_prompt_tokens())
  18. def continuation_prompt_tokens(self, cli=None, width=None):
  19. if width is None:
  20. width = self._width()
  21. return [
  22. (Token.Prompt, (' ' * (width - 5)) + '...: '),
  23. ]
  24. def rewrite_prompt_tokens(self):
  25. width = self._width()
  26. return [
  27. (Token.Prompt, ('-' * (width - 2)) + '> '),
  28. ]
  29. def out_prompt_tokens(self):
  30. return [
  31. (Token.OutPrompt, 'Out['),
  32. (Token.OutPromptNum, str(self.shell.execution_count)),
  33. (Token.OutPrompt, ']: '),
  34. ]
  35. class ClassicPrompts(Prompts):
  36. def in_prompt_tokens(self, cli=None):
  37. return [
  38. (Token.Prompt, '>>> '),
  39. ]
  40. def continuation_prompt_tokens(self, cli=None, width=None):
  41. return [
  42. (Token.Prompt, '... ')
  43. ]
  44. def rewrite_prompt_tokens(self):
  45. return []
  46. def out_prompt_tokens(self):
  47. return []
  48. class RichPromptDisplayHook(DisplayHook):
  49. """Subclass of base display hook using coloured prompt"""
  50. def write_output_prompt(self):
  51. sys.stdout.write(self.shell.separate_out)
  52. # If we're not displaying a prompt, it effectively ends with a newline,
  53. # because the output will be left-aligned.
  54. self.prompt_end_newline = True
  55. if self.do_full_cache:
  56. tokens = self.shell.prompts.out_prompt_tokens()
  57. prompt_txt = ''.join(s for t, s in tokens)
  58. if prompt_txt and not prompt_txt.endswith('\n'):
  59. # Ask for a newline before multiline output
  60. self.prompt_end_newline = False
  61. if self.shell.pt_cli:
  62. self.shell.pt_cli.print_tokens(tokens)
  63. else:
  64. sys.stdout.write(prompt_txt)