106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
import asyncio
|
|
import errno
|
|
import signal
|
|
|
|
from pexpect import EOF
|
|
|
|
@asyncio.coroutine
|
|
def expect_async(expecter, timeout=None):
|
|
# First process data that was previously read - if it maches, we don't need
|
|
# async stuff.
|
|
previously_read = expecter.spawn.buffer
|
|
expecter.spawn._buffer = expecter.spawn.buffer_type()
|
|
expecter.spawn._before = expecter.spawn.buffer_type()
|
|
idx = expecter.new_data(previously_read)
|
|
if idx is not None:
|
|
return idx
|
|
if not expecter.spawn.async_pw_transport:
|
|
pw = PatternWaiter()
|
|
pw.set_expecter(expecter)
|
|
transport, pw = yield from asyncio.get_event_loop()\
|
|
.connect_read_pipe(lambda: pw, expecter.spawn)
|
|
expecter.spawn.async_pw_transport = pw, transport
|
|
else:
|
|
pw, transport = expecter.spawn.async_pw_transport
|
|
pw.set_expecter(expecter)
|
|
transport.resume_reading()
|
|
try:
|
|
return (yield from asyncio.wait_for(pw.fut, timeout))
|
|
except asyncio.TimeoutError as e:
|
|
transport.pause_reading()
|
|
return expecter.timeout(e)
|
|
|
|
@asyncio.coroutine
|
|
def repl_run_command_async(repl, cmdlines, timeout=-1):
|
|
res = []
|
|
repl.child.sendline(cmdlines[0])
|
|
for line in cmdlines[1:]:
|
|
yield from repl._expect_prompt(timeout=timeout, async_=True)
|
|
res.append(repl.child.before)
|
|
repl.child.sendline(line)
|
|
|
|
# Command was fully submitted, now wait for the next prompt
|
|
prompt_idx = yield from repl._expect_prompt(timeout=timeout, async_=True)
|
|
if prompt_idx == 1:
|
|
# We got the continuation prompt - command was incomplete
|
|
repl.child.kill(signal.SIGINT)
|
|
yield from repl._expect_prompt(timeout=1, async_=True)
|
|
raise ValueError("Continuation prompt found - input was incomplete:")
|
|
return u''.join(res + [repl.child.before])
|
|
|
|
class PatternWaiter(asyncio.Protocol):
|
|
transport = None
|
|
|
|
def set_expecter(self, expecter):
|
|
self.expecter = expecter
|
|
self.fut = asyncio.Future()
|
|
|
|
def found(self, result):
|
|
if not self.fut.done():
|
|
self.fut.set_result(result)
|
|
self.transport.pause_reading()
|
|
|
|
def error(self, exc):
|
|
if not self.fut.done():
|
|
self.fut.set_exception(exc)
|
|
self.transport.pause_reading()
|
|
|
|
def connection_made(self, transport):
|
|
self.transport = transport
|
|
|
|
def data_received(self, data):
|
|
spawn = self.expecter.spawn
|
|
s = spawn._decoder.decode(data)
|
|
spawn._log(s, 'read')
|
|
|
|
if self.fut.done():
|
|
spawn._buffer.write(s)
|
|
return
|
|
|
|
try:
|
|
index = self.expecter.new_data(s)
|
|
if index is not None:
|
|
# Found a match
|
|
self.found(index)
|
|
except Exception as e:
|
|
self.expecter.errored()
|
|
self.error(e)
|
|
|
|
def eof_received(self):
|
|
# N.B. If this gets called, async will close the pipe (the spawn object)
|
|
# for us
|
|
try:
|
|
self.expecter.spawn.flag_eof = True
|
|
index = self.expecter.eof()
|
|
except EOF as e:
|
|
self.error(e)
|
|
else:
|
|
self.found(index)
|
|
|
|
def connection_lost(self, exc):
|
|
if isinstance(exc, OSError) and exc.errno == errno.EIO:
|
|
# We may get here without eof_received being called, e.g on Linux
|
|
self.eof_received()
|
|
elif exc is not None:
|
|
self.error(exc)
|