1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
|
# Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
#
# SPDX-License-Identifier: GPL-2.0
# Logic to spawn a sub-process and interact with its stdio.
import os
import re
import pty
import signal
import select
import time
class Timeout(Exception):
"""An exception sub-class that indicates that a timeout occurred."""
pass
class Spawn(object):
"""Represents the stdio of a freshly created sub-process. Commands may be
sent to the process, and responses waited for.
Members:
output: accumulated output from expect()
"""
def __init__(self, args, cwd=None):
"""Spawn (fork/exec) the sub-process.
Args:
args: array of processs arguments. argv[0] is the command to
execute.
cwd: the directory to run the process in, or None for no change.
Returns:
Nothing.
"""
self.waited = False
self.buf = ''
self.output = ''
self.logfile_read = None
self.before = ''
self.after = ''
self.timeout = None
# http://stackoverflow.com/questions/7857352/python-regex-to-match-vt100-escape-sequences
# Note that re.I doesn't seem to work with this regex (or perhaps the
# version of Python in Ubuntu 14.04), hence the inclusion of a-z inside
# [] instead.
self.re_vt100 = re.compile('(\x1b\[|\x9b)[^@-_a-z]*[@-_a-z]|\x1b[@-_a-z]')
(self.pid, self.fd) = pty.fork()
if self.pid == 0:
try:
# For some reason, SIGHUP is set to SIG_IGN at this point when
# run under "go" (www.go.cd). Perhaps this happens under any
# background (non-interactive) system?
signal.signal(signal.SIGHUP, signal.SIG_DFL)
if cwd:
os.chdir(cwd)
os.execvp(args[0], args)
except:
print 'CHILD EXECEPTION:'
import traceback
traceback.print_exc()
finally:
os._exit(255)
try:
self.poll = select.poll()
self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
except:
self.close()
raise
def kill(self, sig):
"""Send unix signal "sig" to the child process.
Args:
sig: The signal number to send.
Returns:
Nothing.
"""
os.kill(self.pid, sig)
def isalive(self):
"""Determine whether the child process is still running.
Args:
None.
Returns:
Boolean indicating whether process is alive.
"""
if self.waited:
return False
w = os.waitpid(self.pid, os.WNOHANG)
if w[0] == 0:
return True
self.waited = True
return False
def send(self, data):
"""Send data to the sub-process's stdin.
Args:
data: The data to send to the process.
Returns:
Nothing.
"""
os.write(self.fd, data)
def expect(self, patterns):
"""Wait for the sub-process to emit specific data.
This function waits for the process to emit one pattern from the
supplied list of patterns, or for a timeout to occur.
Args:
patterns: A list of strings or regex objects that we expect to
see in the sub-process' stdout.
Returns:
The index within the patterns array of the pattern the process
emitted.
Notable exceptions:
Timeout, if the process did not emit any of the patterns within
the expected time.
"""
for pi in xrange(len(patterns)):
if type(patterns[pi]) == type(''):
patterns[pi] = re.compile(patterns[pi])
tstart_s = time.time()
try:
while True:
earliest_m = None
earliest_pi = None
for pi in xrange(len(patterns)):
pattern = patterns[pi]
m = pattern.search(self.buf)
if not m:
continue
if earliest_m and m.start() >= earliest_m.start():
continue
earliest_m = m
earliest_pi = pi
if earliest_m:
pos = earliest_m.start()
posafter = earliest_m.end()
self.before = self.buf[:pos]
self.after = self.buf[pos:posafter]
self.output += self.buf[:posafter]
self.buf = self.buf[posafter:]
return earliest_pi
tnow_s = time.time()
if self.timeout:
tdelta_ms = (tnow_s - tstart_s) * 1000
poll_maxwait = self.timeout - tdelta_ms
if tdelta_ms > self.timeout:
raise Timeout()
else:
poll_maxwait = None
events = self.poll.poll(poll_maxwait)
if not events:
raise Timeout()
c = os.read(self.fd, 1024)
if not c:
raise EOFError()
if self.logfile_read:
self.logfile_read.write(c)
self.buf += c
# count=0 is supposed to be the default, which indicates
# unlimited substitutions, but in practice the version of
# Python in Ubuntu 14.04 appears to default to count=2!
self.buf = self.re_vt100.sub('', self.buf, count=1000000)
finally:
if self.logfile_read:
self.logfile_read.flush()
def close(self):
"""Close the stdio connection to the sub-process.
This also waits a reasonable time for the sub-process to stop running.
Args:
None.
Returns:
Nothing.
"""
os.close(self.fd)
for i in xrange(100):
if not self.isalive():
break
time.sleep(0.1)
def get_expect_output(self):
"""Return the output read by expect()
Returns:
The output processed by expect(), as a string.
"""
return self.output
|
