async.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. import asyncio
  2. import errno
  3. from pexpect import EOF
  4. @asyncio.coroutine
  5. def expect_async(expecter, timeout=None):
  6. # First process data that was previously read - if it maches, we don't need
  7. # async stuff.
  8. previously_read = expecter.spawn.buffer
  9. expecter.spawn.buffer = expecter.spawn.string_type()
  10. idx = expecter.new_data(previously_read)
  11. if idx is not None:
  12. return idx
  13. transport, pw = yield from asyncio.get_event_loop()\
  14. .connect_read_pipe(lambda: PatternWaiter(expecter), expecter.spawn)
  15. try:
  16. return (yield from asyncio.wait_for(pw.fut, timeout))
  17. except asyncio.TimeoutError as e:
  18. transport.pause_reading()
  19. return expecter.timeout(e)
  20. class PatternWaiter(asyncio.Protocol):
  21. transport = None
  22. def __init__(self, expecter):
  23. self.expecter = expecter
  24. self.fut = asyncio.Future()
  25. def found(self, result):
  26. if not self.fut.done():
  27. self.fut.set_result(result)
  28. self.transport.pause_reading()
  29. def error(self, exc):
  30. if not self.fut.done():
  31. self.fut.set_exception(exc)
  32. self.transport.pause_reading()
  33. def connection_made(self, transport):
  34. self.transport = transport
  35. def data_received(self, data):
  36. spawn = self.expecter.spawn
  37. s = spawn._decoder.decode(data)
  38. spawn._log(s, 'read')
  39. if self.fut.done():
  40. spawn.buffer += s
  41. return
  42. try:
  43. index = self.expecter.new_data(s)
  44. if index is not None:
  45. # Found a match
  46. self.found(index)
  47. except Exception as e:
  48. self.expecter.errored()
  49. self.error(e)
  50. def eof_received(self):
  51. # N.B. If this gets called, async will close the pipe (the spawn object)
  52. # for us
  53. try:
  54. self.expecter.spawn.flag_eof = True
  55. index = self.expecter.eof()
  56. except EOF as e:
  57. self.error(e)
  58. else:
  59. self.found(index)
  60. def connection_lost(self, exc):
  61. if isinstance(exc, OSError) and exc.errno == errno.EIO:
  62. # We may get here without eof_received being called, e.g on Linux
  63. self.eof_received()
  64. elif exc is not None:
  65. self.error(exc)