mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 1 | """Convenience functions for use by tests or whomever. |
| 2 | """ |
| 3 | |
mbligh | ff88e4e | 2006-05-25 19:34:52 +0000 | [diff] [blame] | 4 | import os,os.path,shutil,urllib,sys,signal,commands,pickle |
apw | bc2867d | 2006-04-06 18:21:16 +0000 | [diff] [blame] | 5 | from error import * |
mbligh | 4e75b0d | 2006-08-29 15:22:44 +0000 | [diff] [blame] | 6 | import re,string |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 7 | |
| 8 | def grep(pattern, file): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 9 | """This is mainly to fix the return code inversion from grep""" |
mbligh | a975fb6 | 2006-04-22 19:56:25 +0000 | [diff] [blame] | 10 | return not system('grep "' + pattern + '" "' + file + '"') |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 11 | |
| 12 | |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 13 | def difflist(list1, list2): |
| 14 | """returns items in list2 that are not in list1""" |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 15 | diff = []; |
| 16 | for x in list2: |
| 17 | if x not in list1: |
| 18 | diff.append(x) |
| 19 | return diff |
| 20 | |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 21 | |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 22 | def cat_file_to_cmd(file, command): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 23 | """equivalent to 'cat file | command' but knows to use zcat or bzcat if appropriate""" |
mbligh | 712cd14 | 2006-04-22 18:57:50 +0000 | [diff] [blame] | 24 | if not os.path.isfile(file): |
| 25 | raise NameError, 'invalid file %s to cat to command %s' % file, command |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 26 | if file.endswith('.bz2'): |
mbligh | a975fb6 | 2006-04-22 19:56:25 +0000 | [diff] [blame] | 27 | system('bzcat ' + file + ' | ' + command) |
mbligh | 3d91491 | 2006-04-22 17:37:19 +0000 | [diff] [blame] | 28 | elif (file.endswith('.gz') or file.endswith('.tgz')): |
mbligh | a975fb6 | 2006-04-22 19:56:25 +0000 | [diff] [blame] | 29 | system('zcat ' + file + ' | ' + command) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 30 | else: |
mbligh | a975fb6 | 2006-04-22 19:56:25 +0000 | [diff] [blame] | 31 | system('cat ' + file + ' | ' + command) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 32 | |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 33 | |
mbligh | 712cd14 | 2006-04-22 18:57:50 +0000 | [diff] [blame] | 34 | def extract_tarball_to_dir(tarball, dir): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 35 | """ |
| 36 | Extract a tarball to a specified directory name instead of whatever |
| 37 | the top level of a tarball is - useful for versioned directory names, etc |
| 38 | """ |
mbligh | 712cd14 | 2006-04-22 18:57:50 +0000 | [diff] [blame] | 39 | if os.path.exists(dir): |
| 40 | raise NameError, 'target %s already exists' % dir |
| 41 | pwd = os.getcwd() |
| 42 | os.chdir(os.path.dirname(os.path.abspath(dir))) |
| 43 | newdir = extract_tarball(tarball) |
| 44 | os.rename(newdir, dir) |
| 45 | os.chdir(pwd) |
| 46 | |
| 47 | |
mbligh | 712cd14 | 2006-04-22 18:57:50 +0000 | [diff] [blame] | 48 | def extract_tarball(tarball): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 49 | """Returns the first found newly created directory by the tarball extraction""" |
mbligh | 712cd14 | 2006-04-22 18:57:50 +0000 | [diff] [blame] | 50 | oldlist = os.listdir('.') |
| 51 | cat_file_to_cmd(tarball, 'tar xf -') |
| 52 | newlist = os.listdir('.') |
| 53 | newfiles = difflist(oldlist, newlist) # what is new dir ? |
| 54 | new_dir = None |
| 55 | for newfile in newfiles: |
| 56 | if (os.path.isdir(newfile)): |
| 57 | return newfile |
| 58 | raise NameError, "extracting tarball produced no dir" |
| 59 | |
mbligh | cdf02a4 | 2006-04-23 02:22:49 +0000 | [diff] [blame] | 60 | |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 61 | def update_version(srcdir, new_version, install, *args, **dargs): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 62 | """Make sure srcdir is version new_version |
| 63 | |
| 64 | If not, delete it and install() the new version |
| 65 | """ |
mbligh | 7290556 | 2006-05-25 01:30:49 +0000 | [diff] [blame] | 66 | versionfile = srcdir + '/.version' |
mbligh | ff88e4e | 2006-05-25 19:34:52 +0000 | [diff] [blame] | 67 | if os.path.exists(srcdir): |
| 68 | if os.path.exists(versionfile): |
| 69 | old_version = pickle.load(open(versionfile, 'r')) |
| 70 | if (old_version != new_version): |
| 71 | system('rm -rf ' + srcdir) |
| 72 | else: |
mbligh | 7290556 | 2006-05-25 01:30:49 +0000 | [diff] [blame] | 73 | system('rm -rf ' + srcdir) |
| 74 | if not os.path.exists(srcdir): |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 75 | install(*args, **dargs) |
mbligh | 7290556 | 2006-05-25 01:30:49 +0000 | [diff] [blame] | 76 | if os.path.exists(srcdir): |
| 77 | pickle.dump(new_version, open(versionfile, 'w')) |
| 78 | |
| 79 | |
mbligh | ea30c8a | 2006-04-22 22:24:25 +0000 | [diff] [blame] | 80 | def is_url(path): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 81 | """true if path is a url |
| 82 | """ |
| 83 | # should cope with other url types here, but we only handle http and ftp |
mbligh | ea30c8a | 2006-04-22 22:24:25 +0000 | [diff] [blame] | 84 | if (path.startswith('http://')) or (path.startswith('ftp://')): |
mbligh | ea30c8a | 2006-04-22 22:24:25 +0000 | [diff] [blame] | 85 | return 1 |
| 86 | return 0 |
| 87 | |
mbligh | cdf02a4 | 2006-04-23 02:22:49 +0000 | [diff] [blame] | 88 | |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 89 | def get_file(src, dest): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 90 | """get a file, either from url or local""" |
mbligh | 3118661 | 2006-04-22 21:55:56 +0000 | [diff] [blame] | 91 | if (src == dest): # no-op here allows clean overrides in tests |
| 92 | return |
mbligh | ea30c8a | 2006-04-22 22:24:25 +0000 | [diff] [blame] | 93 | if (is_url(src)): |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 94 | print 'PWD: ' + os.getcwd() |
| 95 | print 'Fetching \n\t', src, '\n\t->', dest |
| 96 | try: |
| 97 | urllib.urlretrieve(src, dest) |
| 98 | except IOError: |
| 99 | sys.stderr.write("Unable to retrieve %s (to %s)\n" % (src, dest)) |
| 100 | sys.exit(1) |
| 101 | return dest |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 102 | shutil.copyfile(src, dest) |
| 103 | return dest |
| 104 | |
mbligh | ea30c8a | 2006-04-22 22:24:25 +0000 | [diff] [blame] | 105 | |
mbligh | 8264186 | 2006-04-23 06:21:36 +0000 | [diff] [blame] | 106 | def unmap_url(srcdir, src, destdir = '.'): |
mbligh | ea30c8a | 2006-04-22 22:24:25 +0000 | [diff] [blame] | 107 | if is_url(src): |
| 108 | dest = destdir + '/' + os.path.basename(src) |
| 109 | get_file(src, dest) |
| 110 | return dest |
mbligh | 8264186 | 2006-04-23 06:21:36 +0000 | [diff] [blame] | 111 | else: |
| 112 | return srcdir + '/' + src |
mbligh | ea30c8a | 2006-04-22 22:24:25 +0000 | [diff] [blame] | 113 | |
| 114 | |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 115 | def basename(path): |
| 116 | i = path.rfind('/'); |
| 117 | return path[i+1:] |
| 118 | |
| 119 | |
| 120 | def force_copy(src, dest): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 121 | """Replace dest with a new copy of src, even if it exists""" |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 122 | if os.path.isfile(dest): |
| 123 | os.remove(dest) |
| 124 | return shutil.copyfile(src, dest) |
| 125 | |
| 126 | |
mbligh | cdf02a4 | 2006-04-23 02:22:49 +0000 | [diff] [blame] | 127 | def file_contains_pattern(file, pattern): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 128 | """Return true if file contains the specified egrep pattern""" |
mbligh | cdf02a4 | 2006-04-23 02:22:49 +0000 | [diff] [blame] | 129 | if not os.path.isfile(file): |
| 130 | raise NameError, 'file %s does not exist' % file |
| 131 | return not system('egrep -q ' + pattern + ' ' + file, ignorestatus = 1) |
| 132 | |
| 133 | |
| 134 | def list_grep(list, pattern): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 135 | """True if any item in list matches the specified pattern.""" |
mbligh | cdf02a4 | 2006-04-23 02:22:49 +0000 | [diff] [blame] | 136 | compiled = re.compile(pattern) |
| 137 | for line in list: |
| 138 | match = compiled.search(line) |
| 139 | if (match): |
| 140 | return 1 |
| 141 | return 0 |
| 142 | |
| 143 | |
mbligh | f49d5cf | 2006-04-23 02:24:42 +0000 | [diff] [blame] | 144 | def get_vmlinux(): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 145 | """Return the full path to vmlinux |
| 146 | |
| 147 | Ahem. This is crap. Pray harder. Bad Martin. |
| 148 | """ |
mbligh | 67b5ece | 2006-04-23 02:53:12 +0000 | [diff] [blame] | 149 | vmlinux = '/boot/vmlinux' |
| 150 | if not os.path.isfile(vmlinux): |
| 151 | raise NameError, 'Cannot find vmlinux' |
| 152 | return vmlinux |
mbligh | f49d5cf | 2006-04-23 02:24:42 +0000 | [diff] [blame] | 153 | |
| 154 | |
| 155 | def get_systemmap(): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 156 | """Return the full path to System.map |
| 157 | |
| 158 | Ahem. This is crap. Pray harder. Bad Martin. |
| 159 | """ |
mbligh | 67b5ece | 2006-04-23 02:53:12 +0000 | [diff] [blame] | 160 | map = '/boot/System.map' |
| 161 | if not os.path.isfile(map): |
| 162 | raise NameError, 'Cannot find System.map' |
| 163 | return map |
| 164 | |
| 165 | |
| 166 | def get_modules_dir(): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 167 | """Return the modules dir for the running kernel version""" |
mbligh | 67b5ece | 2006-04-23 02:53:12 +0000 | [diff] [blame] | 168 | kernel_version = system_output('uname -r') |
| 169 | return '/lib/modules/%s/kernel' % kernel_version |
mbligh | f49d5cf | 2006-04-23 02:24:42 +0000 | [diff] [blame] | 170 | |
| 171 | |
mbligh | 5970cf0 | 2006-08-06 15:39:22 +0000 | [diff] [blame] | 172 | def get_cpu_arch(): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 173 | """Work out which CPU architecture we're running on""" |
mbligh | cdf02a4 | 2006-04-23 02:22:49 +0000 | [diff] [blame] | 174 | f = open('/proc/cpuinfo', 'r') |
| 175 | cpuinfo = f.readlines() |
| 176 | f.close() |
| 177 | if list_grep(cpuinfo, '^cpu.*(RS64|POWER3|Broadband Engine)'): |
| 178 | return 'power' |
| 179 | elif list_grep(cpuinfo, '^cpu.*POWER4'): |
| 180 | return 'power4' |
| 181 | elif list_grep(cpuinfo, '^cpu.*POWER5'): |
| 182 | return 'power5' |
| 183 | elif list_grep(cpuinfo, '^cpu.*POWER6'): |
| 184 | return 'power6' |
| 185 | elif list_grep(cpuinfo, '^cpu.*PPC970'): |
| 186 | return 'power970' |
| 187 | elif list_grep(cpuinfo, 'Opteron'): |
| 188 | return 'x86_64' |
| 189 | elif list_grep(cpuinfo, 'GenuineIntel') and list_grep(cpuinfo, '48 bits virtual'): |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 190 | return 'x86_64' |
| 191 | else: |
| 192 | return 'i386' |
| 193 | |
| 194 | |
mbligh | 5970cf0 | 2006-08-06 15:39:22 +0000 | [diff] [blame] | 195 | def get_kernel_arch(): |
| 196 | """Work out the current kernel architecture (as a kernel arch)""" |
| 197 | arch = os.popen('uname -m').read().rstrip() |
| 198 | if ((arch == 'i586') or (arch == 'i686')): |
| 199 | return 'i386' |
mbligh | cdf02a4 | 2006-04-23 02:22:49 +0000 | [diff] [blame] | 200 | else: |
| 201 | return arch |
| 202 | |
| 203 | |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 204 | def kernelexpand(kernel): |
| 205 | # if not (kernel.startswith('http://') or kernel.startswith('ftp://') or os.path.isfile(kernel)): |
mbligh | 5629af0 | 2006-09-15 01:54:23 +0000 | [diff] [blame] | 206 | if kernel.find('/') < 0: # contains no path. |
mbligh | 534015f | 2006-09-15 03:28:56 +0000 | [diff] [blame^] | 207 | autodir = os.environ['AUTODIR'] |
| 208 | kernelexpand = os.path.join(autodir, 'tools/kernelexpand') |
mbligh | 5629af0 | 2006-09-15 01:54:23 +0000 | [diff] [blame] | 209 | w, r = os.popen2(kernelexpand + ' ' + kernel) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 210 | |
| 211 | kernel = r.readline().strip() |
| 212 | r.close() |
| 213 | w.close() |
mbligh | 534015f | 2006-09-15 03:28:56 +0000 | [diff] [blame^] | 214 | return kernel.split() |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 215 | |
| 216 | |
| 217 | def count_cpus(): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 218 | """number of CPUs in the local machine according to /proc/cpuinfo""" |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 219 | f = file('/proc/cpuinfo', 'r') |
| 220 | cpus = 0 |
| 221 | for line in f.readlines(): |
| 222 | if line.startswith('processor'): |
| 223 | cpus += 1 |
| 224 | return cpus |
| 225 | |
mbligh | a975fb6 | 2006-04-22 19:56:25 +0000 | [diff] [blame] | 226 | def system(cmd, ignorestatus = 0): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 227 | """os.system replacement |
| 228 | |
| 229 | We have our own definition of system here, as the stock os.system doesn't |
| 230 | correctly handle sigpipe |
| 231 | (ie things like "yes | head" will hang because yes doesn't get the SIGPIPE). |
| 232 | |
| 233 | Also the stock os.system didn't raise errors based on exit status, this |
| 234 | version does unless you explicitly specify ignorestatus=1 |
| 235 | """ |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 236 | signal.signal(signal.SIGPIPE, signal.SIG_DFL) |
| 237 | try: |
apw | bc2867d | 2006-04-06 18:21:16 +0000 | [diff] [blame] | 238 | status = os.system(cmd) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 239 | finally: |
| 240 | signal.signal(signal.SIGPIPE, signal.SIG_IGN) |
mbligh | a975fb6 | 2006-04-22 19:56:25 +0000 | [diff] [blame] | 241 | |
mbligh | 67b5ece | 2006-04-23 02:53:12 +0000 | [diff] [blame] | 242 | if ((status != 0) and not ignorestatus): |
apw | bc2867d | 2006-04-06 18:21:16 +0000 | [diff] [blame] | 243 | raise CmdError(cmd, status) |
mbligh | a975fb6 | 2006-04-22 19:56:25 +0000 | [diff] [blame] | 244 | return status |
mbligh | 3d91491 | 2006-04-22 17:37:19 +0000 | [diff] [blame] | 245 | |
| 246 | |
mbligh | 67b5ece | 2006-04-23 02:53:12 +0000 | [diff] [blame] | 247 | def system_output(command, ignorestatus = 0): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 248 | """Run command and return its output |
| 249 | |
| 250 | ignorestatus |
| 251 | whether to raise a CmdError if command has a nonzero exit status |
| 252 | """ |
mbligh | 67b5ece | 2006-04-23 02:53:12 +0000 | [diff] [blame] | 253 | (result, data) = commands.getstatusoutput(command) |
| 254 | if ((result != 0) and not ignorestatus): |
| 255 | raise CmdError, 'command failed: ' + command |
| 256 | return data |
| 257 | |
| 258 | |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 259 | def where_art_thy_filehandles(): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 260 | """Dump the current list of filehandles""" |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 261 | os.system("ls -l /proc/%d/fd >> /dev/tty" % os.getpid()) |
| 262 | |
| 263 | |
| 264 | def print_to_tty(string): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 265 | """Output string straight to the tty""" |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 266 | os.system("echo " + string + " >> /dev/tty") |
| 267 | |
| 268 | |
mbligh | b8a14e3 | 2006-05-06 00:17:35 +0000 | [diff] [blame] | 269 | def dump_object(object): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 270 | """Dump an object's attributes and methods |
| 271 | |
| 272 | kind of like dir() |
| 273 | """ |
mbligh | b8a14e3 | 2006-05-06 00:17:35 +0000 | [diff] [blame] | 274 | for item in object.__dict__.iteritems(): |
| 275 | print item |
| 276 | try: |
| 277 | (key,value) = item |
| 278 | dump_object(value) |
| 279 | except: |
| 280 | continue |
| 281 | |
| 282 | |
mbligh | 4b08966 | 2006-06-14 22:34:58 +0000 | [diff] [blame] | 283 | def environ(env_key): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 284 | """return the requested environment variable, or '' if unset""" |
mbligh | 4b08966 | 2006-06-14 22:34:58 +0000 | [diff] [blame] | 285 | if (os.environ.has_key(env_key)): |
| 286 | return os.environ(env_key) |
| 287 | else: |
| 288 | return '' |
| 289 | |
| 290 | |
| 291 | def prepend_path(newpath, oldpath): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 292 | """prepend newpath to oldpath""" |
mbligh | 4b08966 | 2006-06-14 22:34:58 +0000 | [diff] [blame] | 293 | if (oldpath): |
| 294 | return newpath + ':' + oldpath |
| 295 | else: |
| 296 | return newpath |
| 297 | |
| 298 | |
| 299 | def append_path(oldpath, newpath): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 300 | """append newpath to oldpath""" |
mbligh | 4b08966 | 2006-06-14 22:34:58 +0000 | [diff] [blame] | 301 | if (oldpath): |
| 302 | return oldpath + ':' + newpath |
| 303 | else: |
| 304 | return newpath |
| 305 | |
| 306 | |
mbligh | 4e75b0d | 2006-08-29 15:22:44 +0000 | [diff] [blame] | 307 | def avgtime_print(dir): |
| 308 | """ Calculate some benchmarking statistics. |
| 309 | Input is a directory containing a file called 'time'. |
| 310 | File contains one-per-line results of /usr/bin/time. |
| 311 | Output is average Elapsed, User, and System time in seconds, |
| 312 | and average CPU percentage. |
| 313 | """ |
| 314 | f = open(dir + "/time") |
| 315 | user = system = elapsed = cpu = count = 0 |
| 316 | r = re.compile('([\d\.]*)user ([\d\.]*)system (\d*):([\d\.]*)elapsed (\d*)%CPU') |
| 317 | for line in f.readlines(): |
| 318 | try: |
| 319 | s = r.match(line); |
mbligh | e1ee267 | 2006-08-29 16:27:15 +0000 | [diff] [blame] | 320 | user += float(s.group(1)) |
| 321 | system += float(s.group(2)) |
| 322 | elapsed += (float(s.group(3)) * 60) + float(s.group(4)) |
| 323 | cpu += float(s.group(5)) |
mbligh | 4e75b0d | 2006-08-29 15:22:44 +0000 | [diff] [blame] | 324 | count += 1 |
| 325 | except: |
| 326 | raise ValueError("badly formatted times") |
| 327 | |
| 328 | f.close() |
| 329 | return "Elapsed: %0.2fs User: %0.2fs System: %0.2fs CPU: %0.0f%%" % \ |
| 330 | (elapsed/count, user/count, system/count, cpu/count) |
| 331 | |
| 332 | |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 333 | class fd_stack: |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 334 | """a stack of fd redirects |
| 335 | |
| 336 | Redirects cause existing fd's to be pushed on the stack; restore() |
| 337 | causes the current set of redirects to be popped, restoring the previous |
| 338 | filehandle destinations. |
| 339 | |
| 340 | Note that we need to redirect both the sys.stdout type descriptor |
| 341 | (which print, etc use) and the low level OS numbered descriptor |
| 342 | which os.system() etc use. |
| 343 | """ |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 344 | |
| 345 | def __init__(self, fd, filehandle): |
| 346 | self.fd = fd # eg 1 |
| 347 | self.filehandle = filehandle # eg sys.stdout |
| 348 | self.stack = [(fd, filehandle)] |
| 349 | |
| 350 | |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 351 | def update_handle(self, new): |
| 352 | if (self.filehandle == sys.stdout): |
| 353 | sys.stdout = new |
| 354 | if (self.filehandle == sys.stderr): |
| 355 | sys.stderr = new |
| 356 | self.filehandle = new |
| 357 | |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 358 | def redirect(self, filename): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 359 | """Redirect output to the specified file |
| 360 | |
| 361 | Overwrites the previous contents, if any. |
| 362 | """ |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 363 | self.filehandle.flush() |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 364 | fdcopy = os.dup(self.fd) |
mbligh | d7341d5 | 2006-09-13 20:02:34 +0000 | [diff] [blame] | 365 | self.stack.append( (fdcopy, self.filehandle, 0) ) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 366 | # self.filehandle = file(filename, 'w') |
| 367 | if (os.path.isfile(filename)): |
| 368 | newfd = os.open(filename, os.O_WRONLY) |
| 369 | else: |
| 370 | newfd = os.open(filename, os.O_WRONLY | os.O_CREAT) |
| 371 | os.dup2(newfd, self.fd) |
| 372 | os.close(newfd) |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 373 | self.update_handle(os.fdopen(self.fd, 'w')) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 374 | |
| 375 | |
| 376 | def tee_redirect(self, filename): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 377 | """Tee output to the specified file |
| 378 | |
| 379 | Overwrites the previous contents, if any. |
| 380 | """ |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 381 | self.filehandle.flush() |
mbligh | 5156634 | 2006-08-29 04:55:40 +0000 | [diff] [blame] | 382 | #print_to_tty("tee_redirect to " + filename) |
| 383 | #where_art_thy_filehandles() |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 384 | fdcopy = os.dup(self.fd) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 385 | r, w = os.pipe() |
| 386 | pid = os.fork() |
| 387 | if pid: # parent |
| 388 | os.close(r) |
| 389 | os.dup2(w, self.fd) |
| 390 | os.close(w) |
| 391 | else: # child |
| 392 | os.close(w) |
| 393 | os.dup2(r, 0) |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 394 | os.dup2(fdcopy, 1) |
| 395 | os.close(r) |
| 396 | os.close(fdcopy) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 397 | os.execlp('tee', 'tee', filename) |
mbligh | d7341d5 | 2006-09-13 20:02:34 +0000 | [diff] [blame] | 398 | self.stack.append( (fdcopy, self.filehandle, pid) ) |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 399 | self.update_handle(os.fdopen(self.fd, 'w')) |
mbligh | 5156634 | 2006-08-29 04:55:40 +0000 | [diff] [blame] | 400 | #where_art_thy_filehandles() |
| 401 | #print_to_tty("done tee_redirect to " + filename) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 402 | |
| 403 | |
| 404 | def restore(self): |
mbligh | c86b0b4 | 2006-07-28 17:35:28 +0000 | [diff] [blame] | 405 | """unredirect one level""" |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 406 | self.filehandle.flush() |
mbligh | cf31533 | 2006-04-02 19:58:30 +0000 | [diff] [blame] | 407 | # print_to_tty("ENTERING RESTORE %d" % self.fd) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 408 | # where_art_thy_filehandles() |
mbligh | d7341d5 | 2006-09-13 20:02:34 +0000 | [diff] [blame] | 409 | (old_fd, old_filehandle, pid) = self.stack.pop() |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 410 | # print_to_tty("old_fd %d" % old_fd) |
| 411 | # print_to_tty("self.fd %d" % self.fd) |
| 412 | self.filehandle.close() # seems to close old_fd as well. |
mbligh | d7341d5 | 2006-09-13 20:02:34 +0000 | [diff] [blame] | 413 | if pid: |
| 414 | os.waitpid(pid, 0) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 415 | # where_art_thy_filehandles() |
| 416 | os.dup2(old_fd, self.fd) |
| 417 | # print_to_tty("CLOSING FD %d" % old_fd) |
| 418 | os.close(old_fd) |
| 419 | # where_art_thy_filehandles() |
mbligh | 78e1702 | 2006-09-13 19:58:12 +0000 | [diff] [blame] | 420 | self.update_handle(old_filehandle) |
mbligh | f4c3532 | 2006-03-13 01:01:10 +0000 | [diff] [blame] | 421 | # where_art_thy_filehandles() |
| 422 | # print_to_tty("EXIT RESTORE %d" % self.fd) |