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