xref: /openbmc/u-boot/test/py/u_boot_spawn.py (revision ad3098f7) 
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
7import os
8import re
9import pty
10import signal
11import select
12import time
13
14class Timeout(Exception):
15    """An exception sub-class that indicates that a timeout occurred."""
16    pass
17
18class Spawn(object):
19    """Represents the stdio of a freshly created sub-process. Commands may be
20    sent to the process, and responses waited for.
21    """
22
23    def __init__(self, args, cwd=None):
24        """Spawn (fork/exec) the sub-process.
25
26        Args:
27            args: array of processs arguments. argv[0] is the command to
28              execute.
29            cwd: the directory to run the process in, or None for no change.
30
31        Returns:
32            Nothing.
33        """
34
35        self.waited = False
36        self.buf = ''
37        self.logfile_read = None
38        self.before = ''
39        self.after = ''
40        self.timeout = None
41
42        (self.pid, self.fd) = pty.fork()
43        if self.pid == 0:
44            try:
45                # For some reason, SIGHUP is set to SIG_IGN at this point when
46                # run under "go" (www.go.cd). Perhaps this happens under any
47                # background (non-interactive) system?
48                signal.signal(signal.SIGHUP, signal.SIG_DFL)
49                if cwd:
50                    os.chdir(cwd)
51                os.execvp(args[0], args)
52            except:
53                print 'CHILD EXECEPTION:'
54                import traceback
55                traceback.print_exc()
56            finally:
57                os._exit(255)
58
59        try:
60            self.poll = select.poll()
61            self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
62        except:
63            self.close()
64            raise
65
66    def kill(self, sig):
67        """Send unix signal "sig" to the child process.
68
69        Args:
70            sig: The signal number to send.
71
72        Returns:
73            Nothing.
74        """
75
76        os.kill(self.pid, sig)
77
78    def isalive(self):
79        """Determine whether the child process is still running.
80
81        Args:
82            None.
83
84        Returns:
85            Boolean indicating whether process is alive.
86        """
87
88        if self.waited:
89            return False
90
91        w = os.waitpid(self.pid, os.WNOHANG)
92        if w[0] == 0:
93            return True
94
95        self.waited = True
96        return False
97
98    def send(self, data):
99        """Send data to the sub-process's stdin.
100
101        Args:
102            data: The data to send to the process.
103
104        Returns:
105            Nothing.
106        """
107
108        os.write(self.fd, data)
109
110    def expect(self, patterns):
111        """Wait for the sub-process to emit specific data.
112
113        This function waits for the process to emit one pattern from the
114        supplied list of patterns, or for a timeout to occur.
115
116        Args:
117            patterns: A list of strings or regex objects that we expect to
118                see in the sub-process' stdout.
119
120        Returns:
121            The index within the patterns array of the pattern the process
122            emitted.
123
124        Notable exceptions:
125            Timeout, if the process did not emit any of the patterns within
126            the expected time.
127        """
128
129        for pi in xrange(len(patterns)):
130            if type(patterns[pi]) == type(''):
131                patterns[pi] = re.compile(patterns[pi])
132
133        tstart_s = time.time()
134        try:
135            while True:
136                earliest_m = None
137                earliest_pi = None
138                for pi in xrange(len(patterns)):
139                    pattern = patterns[pi]
140                    m = pattern.search(self.buf)
141                    if not m:
142                        continue
143                    if earliest_m and m.start() >= earliest_m.start():
144                        continue
145                    earliest_m = m
146                    earliest_pi = pi
147                if earliest_m:
148                    pos = earliest_m.start()
149                    posafter = earliest_m.end()
150                    self.before = self.buf[:pos]
151                    self.after = self.buf[pos:posafter]
152                    self.buf = self.buf[posafter:]
153                    return earliest_pi
154                tnow_s = time.time()
155                if self.timeout:
156                    tdelta_ms = (tnow_s - tstart_s) * 1000
157                    poll_maxwait = self.timeout - tdelta_ms
158                    if tdelta_ms > self.timeout:
159                        raise Timeout()
160                else:
161                    poll_maxwait = None
162                events = self.poll.poll(poll_maxwait)
163                if not events:
164                    raise Timeout()
165                c = os.read(self.fd, 1024)
166                if not c:
167                    raise EOFError()
168                if self.logfile_read:
169                    self.logfile_read.write(c)
170                self.buf += c
171        finally:
172            if self.logfile_read:
173                self.logfile_read.flush()
174
175    def close(self):
176        """Close the stdio connection to the sub-process.
177
178        This also waits a reasonable time for the sub-process to stop running.
179
180        Args:
181            None.
182
183        Returns:
184            Nothing.
185        """
186
187        os.close(self.fd)
188        for i in xrange(100):
189            if not self.isalive():
190                break
191            time.sleep(0.1)
192