expect.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300
  1. import time
  2. from .exceptions import EOF, TIMEOUT
  3. class Expecter(object):
  4. def __init__(self, spawn, searcher, searchwindowsize=-1):
  5. self.spawn = spawn
  6. self.searcher = searcher
  7. if searchwindowsize == -1:
  8. searchwindowsize = spawn.searchwindowsize
  9. self.searchwindowsize = searchwindowsize
  10. def new_data(self, data):
  11. spawn = self.spawn
  12. searcher = self.searcher
  13. incoming = spawn.buffer + data
  14. freshlen = len(data)
  15. index = searcher.search(incoming, freshlen, self.searchwindowsize)
  16. if index >= 0:
  17. spawn.buffer = incoming[searcher.end:]
  18. spawn.before = incoming[: searcher.start]
  19. spawn.after = incoming[searcher.start: searcher.end]
  20. spawn.match = searcher.match
  21. spawn.match_index = index
  22. # Found a match
  23. return index
  24. spawn.buffer = incoming
  25. def eof(self, err=None):
  26. spawn = self.spawn
  27. from . import EOF
  28. spawn.before = spawn.buffer
  29. spawn.buffer = spawn.string_type()
  30. spawn.after = EOF
  31. index = self.searcher.eof_index
  32. if index >= 0:
  33. spawn.match = EOF
  34. spawn.match_index = index
  35. return index
  36. else:
  37. spawn.match = None
  38. spawn.match_index = None
  39. msg = str(spawn)
  40. msg += '\nsearcher: %s' % self.searcher
  41. if err is not None:
  42. msg = str(err) + '\n' + msg
  43. raise EOF(msg)
  44. def timeout(self, err=None):
  45. spawn = self.spawn
  46. from . import TIMEOUT
  47. spawn.before = spawn.buffer
  48. spawn.after = TIMEOUT
  49. index = self.searcher.timeout_index
  50. if index >= 0:
  51. spawn.match = TIMEOUT
  52. spawn.match_index = index
  53. return index
  54. else:
  55. spawn.match = None
  56. spawn.match_index = None
  57. msg = str(spawn)
  58. msg += '\nsearcher: %s' % self.searcher
  59. if err is not None:
  60. msg = str(err) + '\n' + msg
  61. raise TIMEOUT(msg)
  62. def errored(self):
  63. spawn = self.spawn
  64. spawn.before = spawn.buffer
  65. spawn.after = None
  66. spawn.match = None
  67. spawn.match_index = None
  68. def expect_loop(self, timeout=-1):
  69. """Blocking expect"""
  70. spawn = self.spawn
  71. from . import EOF, TIMEOUT
  72. if timeout is not None:
  73. end_time = time.time() + timeout
  74. try:
  75. incoming = spawn.buffer
  76. spawn.buffer = spawn.string_type() # Treat buffer as new data
  77. while True:
  78. idx = self.new_data(incoming)
  79. # Keep reading until exception or return.
  80. if idx is not None:
  81. return idx
  82. # No match at this point
  83. if (timeout is not None) and (timeout < 0):
  84. return self.timeout()
  85. # Still have time left, so read more data
  86. incoming = spawn.read_nonblocking(spawn.maxread, timeout)
  87. if self.spawn.delayafterread is not None:
  88. time.sleep(self.spawn.delayafterread)
  89. if timeout is not None:
  90. timeout = end_time - time.time()
  91. except EOF as e:
  92. return self.eof(e)
  93. except TIMEOUT as e:
  94. return self.timeout(e)
  95. except:
  96. self.errored()
  97. raise
  98. class searcher_string(object):
  99. '''This is a plain string search helper for the spawn.expect_any() method.
  100. This helper class is for speed. For more powerful regex patterns
  101. see the helper class, searcher_re.
  102. Attributes:
  103. eof_index - index of EOF, or -1
  104. timeout_index - index of TIMEOUT, or -1
  105. After a successful match by the search() method the following attributes
  106. are available:
  107. start - index into the buffer, first byte of match
  108. end - index into the buffer, first byte after match
  109. match - the matching string itself
  110. '''
  111. def __init__(self, strings):
  112. '''This creates an instance of searcher_string. This argument 'strings'
  113. may be a list; a sequence of strings; or the EOF or TIMEOUT types. '''
  114. self.eof_index = -1
  115. self.timeout_index = -1
  116. self._strings = []
  117. for n, s in enumerate(strings):
  118. if s is EOF:
  119. self.eof_index = n
  120. continue
  121. if s is TIMEOUT:
  122. self.timeout_index = n
  123. continue
  124. self._strings.append((n, s))
  125. def __str__(self):
  126. '''This returns a human-readable string that represents the state of
  127. the object.'''
  128. ss = [(ns[0], ' %d: "%s"' % ns) for ns in self._strings]
  129. ss.append((-1, 'searcher_string:'))
  130. if self.eof_index >= 0:
  131. ss.append((self.eof_index, ' %d: EOF' % self.eof_index))
  132. if self.timeout_index >= 0:
  133. ss.append((self.timeout_index,
  134. ' %d: TIMEOUT' % self.timeout_index))
  135. ss.sort()
  136. ss = list(zip(*ss))[1]
  137. return '\n'.join(ss)
  138. def search(self, buffer, freshlen, searchwindowsize=None):
  139. '''This searches 'buffer' for the first occurence of one of the search
  140. strings. 'freshlen' must indicate the number of bytes at the end of
  141. 'buffer' which have not been searched before. It helps to avoid
  142. searching the same, possibly big, buffer over and over again.
  143. See class spawn for the 'searchwindowsize' argument.
  144. If there is a match this returns the index of that string, and sets
  145. 'start', 'end' and 'match'. Otherwise, this returns -1. '''
  146. first_match = None
  147. # 'freshlen' helps a lot here. Further optimizations could
  148. # possibly include:
  149. #
  150. # using something like the Boyer-Moore Fast String Searching
  151. # Algorithm; pre-compiling the search through a list of
  152. # strings into something that can scan the input once to
  153. # search for all N strings; realize that if we search for
  154. # ['bar', 'baz'] and the input is '...foo' we need not bother
  155. # rescanning until we've read three more bytes.
  156. #
  157. # Sadly, I don't know enough about this interesting topic. /grahn
  158. for index, s in self._strings:
  159. if searchwindowsize is None:
  160. # the match, if any, can only be in the fresh data,
  161. # or at the very end of the old data
  162. offset = -(freshlen + len(s))
  163. else:
  164. # better obey searchwindowsize
  165. offset = -searchwindowsize
  166. n = buffer.find(s, offset)
  167. if n >= 0 and (first_match is None or n < first_match):
  168. first_match = n
  169. best_index, best_match = index, s
  170. if first_match is None:
  171. return -1
  172. self.match = best_match
  173. self.start = first_match
  174. self.end = self.start + len(self.match)
  175. return best_index
  176. class searcher_re(object):
  177. '''This is regular expression string search helper for the
  178. spawn.expect_any() method. This helper class is for powerful
  179. pattern matching. For speed, see the helper class, searcher_string.
  180. Attributes:
  181. eof_index - index of EOF, or -1
  182. timeout_index - index of TIMEOUT, or -1
  183. After a successful match by the search() method the following attributes
  184. are available:
  185. start - index into the buffer, first byte of match
  186. end - index into the buffer, first byte after match
  187. match - the re.match object returned by a succesful re.search
  188. '''
  189. def __init__(self, patterns):
  190. '''This creates an instance that searches for 'patterns' Where
  191. 'patterns' may be a list or other sequence of compiled regular
  192. expressions, or the EOF or TIMEOUT types.'''
  193. self.eof_index = -1
  194. self.timeout_index = -1
  195. self._searches = []
  196. for n, s in zip(list(range(len(patterns))), patterns):
  197. if s is EOF:
  198. self.eof_index = n
  199. continue
  200. if s is TIMEOUT:
  201. self.timeout_index = n
  202. continue
  203. self._searches.append((n, s))
  204. def __str__(self):
  205. '''This returns a human-readable string that represents the state of
  206. the object.'''
  207. #ss = [(n, ' %d: re.compile("%s")' %
  208. # (n, repr(s.pattern))) for n, s in self._searches]
  209. ss = list()
  210. for n, s in self._searches:
  211. try:
  212. ss.append((n, ' %d: re.compile("%s")' % (n, s.pattern)))
  213. except UnicodeEncodeError:
  214. # for test cases that display __str__ of searches, dont throw
  215. # another exception just because stdout is ascii-only, using
  216. # repr()
  217. ss.append((n, ' %d: re.compile(%r)' % (n, s.pattern)))
  218. ss.append((-1, 'searcher_re:'))
  219. if self.eof_index >= 0:
  220. ss.append((self.eof_index, ' %d: EOF' % self.eof_index))
  221. if self.timeout_index >= 0:
  222. ss.append((self.timeout_index, ' %d: TIMEOUT' %
  223. self.timeout_index))
  224. ss.sort()
  225. ss = list(zip(*ss))[1]
  226. return '\n'.join(ss)
  227. def search(self, buffer, freshlen, searchwindowsize=None):
  228. '''This searches 'buffer' for the first occurence of one of the regular
  229. expressions. 'freshlen' must indicate the number of bytes at the end of
  230. 'buffer' which have not been searched before.
  231. See class spawn for the 'searchwindowsize' argument.
  232. If there is a match this returns the index of that string, and sets
  233. 'start', 'end' and 'match'. Otherwise, returns -1.'''
  234. first_match = None
  235. # 'freshlen' doesn't help here -- we cannot predict the
  236. # length of a match, and the re module provides no help.
  237. if searchwindowsize is None:
  238. searchstart = 0
  239. else:
  240. searchstart = max(0, len(buffer) - searchwindowsize)
  241. for index, s in self._searches:
  242. match = s.search(buffer, searchstart)
  243. if match is None:
  244. continue
  245. n = match.start()
  246. if first_match is None or n < first_match:
  247. first_match = n
  248. the_match = match
  249. best_index = index
  250. if first_match is None:
  251. return -1
  252. self.start = first_match
  253. self.match = the_match
  254. self.end = self.match.end()
  255. return best_index