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