You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

122 lines
5.0 KiB

4 years ago
  1. """Generic wrapper for read-eval-print-loops, a.k.a. interactive shells
  2. """
  3. import os.path
  4. import signal
  5. import sys
  6. import pexpect
  7. PY3 = (sys.version_info[0] >= 3)
  8. if PY3:
  9. basestring = str
  10. PEXPECT_PROMPT = u'[PEXPECT_PROMPT>'
  11. PEXPECT_CONTINUATION_PROMPT = u'[PEXPECT_PROMPT+'
  12. class REPLWrapper(object):
  13. """Wrapper for a REPL.
  14. :param cmd_or_spawn: This can either be an instance of :class:`pexpect.spawn`
  15. in which a REPL has already been started, or a str command to start a new
  16. REPL process.
  17. :param str orig_prompt: The prompt to expect at first.
  18. :param str prompt_change: A command to change the prompt to something more
  19. unique. If this is ``None``, the prompt will not be changed. This will
  20. be formatted with the new and continuation prompts as positional
  21. parameters, so you can use ``{}`` style formatting to insert them into
  22. the command.
  23. :param str new_prompt: The more unique prompt to expect after the change.
  24. :param str extra_init_cmd: Commands to do extra initialisation, such as
  25. disabling pagers.
  26. """
  27. def __init__(self, cmd_or_spawn, orig_prompt, prompt_change,
  28. new_prompt=PEXPECT_PROMPT,
  29. continuation_prompt=PEXPECT_CONTINUATION_PROMPT,
  30. extra_init_cmd=None):
  31. if isinstance(cmd_or_spawn, basestring):
  32. self.child = pexpect.spawn(cmd_or_spawn, echo=False, encoding='utf-8')
  33. else:
  34. self.child = cmd_or_spawn
  35. if self.child.echo:
  36. # Existing spawn instance has echo enabled, disable it
  37. # to prevent our input from being repeated to output.
  38. self.child.setecho(False)
  39. self.child.waitnoecho()
  40. if prompt_change is None:
  41. self.prompt = orig_prompt
  42. else:
  43. self.set_prompt(orig_prompt,
  44. prompt_change.format(new_prompt, continuation_prompt))
  45. self.prompt = new_prompt
  46. self.continuation_prompt = continuation_prompt
  47. self._expect_prompt()
  48. if extra_init_cmd is not None:
  49. self.run_command(extra_init_cmd)
  50. def set_prompt(self, orig_prompt, prompt_change):
  51. self.child.expect(orig_prompt)
  52. self.child.sendline(prompt_change)
  53. def _expect_prompt(self, timeout=-1):
  54. return self.child.expect_exact([self.prompt, self.continuation_prompt],
  55. timeout=timeout)
  56. def run_command(self, command, timeout=-1):
  57. """Send a command to the REPL, wait for and return output.
  58. :param str command: The command to send. Trailing newlines are not needed.
  59. This should be a complete block of input that will trigger execution;
  60. if a continuation prompt is found after sending input, :exc:`ValueError`
  61. will be raised.
  62. :param int timeout: How long to wait for the next prompt. -1 means the
  63. default from the :class:`pexpect.spawn` object (default 30 seconds).
  64. None means to wait indefinitely.
  65. """
  66. # Split up multiline commands and feed them in bit-by-bit
  67. cmdlines = command.splitlines()
  68. # splitlines ignores trailing newlines - add it back in manually
  69. if command.endswith('\n'):
  70. cmdlines.append('')
  71. if not cmdlines:
  72. raise ValueError("No command was given")
  73. res = []
  74. self.child.sendline(cmdlines[0])
  75. for line in cmdlines[1:]:
  76. self._expect_prompt(timeout=timeout)
  77. res.append(self.child.before)
  78. self.child.sendline(line)
  79. # Command was fully submitted, now wait for the next prompt
  80. if self._expect_prompt(timeout=timeout) == 1:
  81. # We got the continuation prompt - command was incomplete
  82. self.child.kill(signal.SIGINT)
  83. self._expect_prompt(timeout=1)
  84. raise ValueError("Continuation prompt found - input was incomplete:\n"
  85. + command)
  86. return u''.join(res + [self.child.before])
  87. def python(command="python"):
  88. """Start a Python shell and return a :class:`REPLWrapper` object."""
  89. return REPLWrapper(command, u">>> ", u"import sys; sys.ps1={0!r}; sys.ps2={1!r}")
  90. def bash(command="bash"):
  91. """Start a bash shell and return a :class:`REPLWrapper` object."""
  92. bashrc = os.path.join(os.path.dirname(__file__), 'bashrc.sh')
  93. child = pexpect.spawn(command, ['--rcfile', bashrc], echo=False,
  94. encoding='utf-8')
  95. # If the user runs 'env', the value of PS1 will be in the output. To avoid
  96. # replwrap seeing that as the next prompt, we'll embed the marker characters
  97. # for invisible characters in the prompt; these show up when inspecting the
  98. # environment variable, but not when bash displays the prompt.
  99. ps1 = PEXPECT_PROMPT[:5] + u'\\[\\]' + PEXPECT_PROMPT[5:]
  100. ps2 = PEXPECT_CONTINUATION_PROMPT[:5] + u'\\[\\]' + PEXPECT_CONTINUATION_PROMPT[5:]
  101. prompt_change = u"PS1='{0}' PS2='{1}' PROMPT_COMMAND=''".format(ps1, ps2)
  102. return REPLWrapper(child, u'\\$', prompt_change,
  103. extra_init_cmd="export PAGER=cat")