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