base_utils.py revision 264cd8f4889cea73fa5fb7873e3243c1c770a1bc
1"""Convenience functions for use by tests or whomever. 2""" 3 4import os,os.path,shutil,urllib,sys,signal,commands,pickle 5from error import * 6import re,string 7 8def grep(pattern, file): 9 """ 10 This is mainly to fix the return code inversion from grep 11 Also handles compressed files. 12 13 returns 1 if the pattern is present in the file, 0 if not. 14 """ 15 command = 'grep "%s" > /dev/null' % pattern 16 ret = cat_file_to_cmd(file, command, ignorestatus = 1) 17 return not ret 18 19 20def difflist(list1, list2): 21 """returns items in list2 that are not in list1""" 22 diff = []; 23 for x in list2: 24 if x not in list1: 25 diff.append(x) 26 return diff 27 28 29def cat_file_to_cmd(file, command, ignorestatus = 0): 30 """ 31 equivalent to 'cat file | command' but knows to use 32 zcat or bzcat if appropriate 33 """ 34 if not os.path.isfile(file): 35 raise NameError, 'invalid file %s to cat to command %s' % file, command 36 if file.endswith('.bz2'): 37 return system('bzcat ' + file + ' | ' + command, ignorestatus) 38 elif (file.endswith('.gz') or file.endswith('.tgz')): 39 return system('zcat ' + file + ' | ' + command, ignorestatus) 40 else: 41 return system('cat ' + file + ' | ' + command, ignorestatus) 42 43 44def extract_tarball_to_dir(tarball, dir): 45 """ 46 Extract a tarball to a specified directory name instead of whatever 47 the top level of a tarball is - useful for versioned directory names, etc 48 """ 49 if os.path.exists(dir): 50 raise NameError, 'target %s already exists' % dir 51 pwd = os.getcwd() 52 os.chdir(os.path.dirname(os.path.abspath(dir))) 53 newdir = extract_tarball(tarball) 54 os.rename(newdir, dir) 55 os.chdir(pwd) 56 57 58def extract_tarball(tarball): 59 """Returns the first found newly created directory by the tarball extraction""" 60 oldlist = os.listdir('.') 61 cat_file_to_cmd(tarball, 'tar xf -') 62 newlist = os.listdir('.') 63 newfiles = difflist(oldlist, newlist) # what is new dir ? 64 new_dir = None 65 for newfile in newfiles: 66 if (os.path.isdir(newfile)): 67 return newfile 68 raise NameError, "extracting tarball produced no dir" 69 70 71def update_version(srcdir, new_version, install, *args, **dargs): 72 """Make sure srcdir is version new_version 73 74 If not, delete it and install() the new version 75 """ 76 versionfile = srcdir + '/.version' 77 if os.path.exists(srcdir): 78 if os.path.exists(versionfile): 79 old_version = pickle.load(open(versionfile, 'r')) 80 if (old_version != new_version): 81 system('rm -rf ' + srcdir) 82 else: 83 system('rm -rf ' + srcdir) 84 if not os.path.exists(srcdir): 85 install(*args, **dargs) 86 if os.path.exists(srcdir): 87 pickle.dump(new_version, open(versionfile, 'w')) 88 89 90def is_url(path): 91 """true if path is a url 92 """ 93 # should cope with other url types here, but we only handle http and ftp 94 if (path.startswith('http://')) or (path.startswith('ftp://')): 95 return 1 96 return 0 97 98 99def get_file(src, dest): 100 """get a file, either from url or local""" 101 if (src == dest): # no-op here allows clean overrides in tests 102 return 103 if (is_url(src)): 104 print 'PWD: ' + os.getcwd() 105 print 'Fetching \n\t', src, '\n\t->', dest 106 try: 107 urllib.urlretrieve(src, dest) 108 except IOError: 109 sys.stderr.write("Unable to retrieve %s (to %s)\n" % (src, dest)) 110 sys.exit(1) 111 return dest 112 shutil.copyfile(src, dest) 113 return dest 114 115 116def unmap_url(srcdir, src, destdir = '.'): 117 """ 118 Receives either a path to a local file or a URL. 119 returns either the path to the local file, or the fetched URL 120 121 unmap_url('/usr/src', 'foo.tar', '/tmp') 122 = '/usr/src/foo.tar' 123 unmap_url('/usr/src', 'http://site/file', '/tmp') 124 = '/tmp/file' 125 (after retrieving it) 126 """ 127 if is_url(src): 128 dest = destdir + '/' + os.path.basename(src) 129 get_file(src, dest) 130 return dest 131 else: 132 return srcdir + '/' + src 133 134 135def basename(path): 136 i = path.rfind('/'); 137 return path[i+1:] 138 139 140def force_copy(src, dest): 141 """Replace dest with a new copy of src, even if it exists""" 142 if os.path.isfile(dest): 143 os.remove(dest) 144 if os.path.isdir(dest): 145 dest = os.path.join(dest, os.path.basename(src)) 146 return shutil.copyfile(src, dest) 147 148 149def force_link(src, dest): 150 """Link src to dest, overwriting it if it exists""" 151 return system("ln -sf %s %s" % (src, dest)) 152 153 154def file_contains_pattern(file, pattern): 155 """Return true if file contains the specified egrep pattern""" 156 if not os.path.isfile(file): 157 raise NameError, 'file %s does not exist' % file 158 return not system('egrep -q "' + pattern + '" ' + file, ignorestatus = 1) 159 160 161def list_grep(list, pattern): 162 """True if any item in list matches the specified pattern.""" 163 compiled = re.compile(pattern) 164 for line in list: 165 match = compiled.search(line) 166 if (match): 167 return 1 168 return 0 169 170def get_os_vendor(): 171 """Try to guess what's the os vendor 172 """ 173 issue = '/etc/issue' 174 175 if not os.path.isfile(issue): 176 return 'Unknown' 177 178 if file_contains_pattern(issue, 'Red Hat'): 179 return 'Red Hat' 180 elif file_contains_pattern(issue, 'Fedora Core'): 181 return 'Fedora Core' 182 elif file_contains_pattern(issue, 'SUSE'): 183 return 'SUSE' 184 elif file_contains_pattern(issue, 'Ubuntu'): 185 return 'Ubuntu' 186 else: 187 return 'Unknown' 188 189 190def get_vmlinux(): 191 """Return the full path to vmlinux 192 193 Ahem. This is crap. Pray harder. Bad Martin. 194 """ 195 vmlinux = '/boot/vmlinux-%s' % system_output('uname -r') 196 if os.path.isfile(vmlinux): 197 return vmlinux 198 vmlinux = '/lib/modules/%s/build/vmlinux' % system_output('uname -r') 199 if os.path.isfile(vmlinux): 200 return vmlinux 201 return None 202 203 204def get_systemmap(): 205 """Return the full path to System.map 206 207 Ahem. This is crap. Pray harder. Bad Martin. 208 """ 209 map = '/boot/System.map-%s' % system_output('uname -r') 210 if os.path.isfile(map): 211 return map 212 map = '/lib/modules/%s/build/System.map' % system_output('uname -r') 213 if os.path.isfile(map): 214 return map 215 return None 216 217 218def get_modules_dir(): 219 """Return the modules dir for the running kernel version""" 220 kernel_version = system_output('uname -r') 221 return '/lib/modules/%s/kernel' % kernel_version 222 223 224def get_cpu_arch(): 225 """Work out which CPU architecture we're running on""" 226 f = open('/proc/cpuinfo', 'r') 227 cpuinfo = f.readlines() 228 f.close() 229 if list_grep(cpuinfo, '^cpu.*(RS64|POWER3|Broadband Engine)'): 230 return 'power' 231 elif list_grep(cpuinfo, '^cpu.*POWER4'): 232 return 'power4' 233 elif list_grep(cpuinfo, '^cpu.*POWER5'): 234 return 'power5' 235 elif list_grep(cpuinfo, '^cpu.*POWER6'): 236 return 'power6' 237 elif list_grep(cpuinfo, '^cpu.*PPC970'): 238 return 'power970' 239 elif list_grep(cpuinfo, 'Opteron'): 240 return 'x86_64' 241 elif list_grep(cpuinfo, 'GenuineIntel') and list_grep(cpuinfo, '48 bits virtual'): 242 return 'x86_64' 243 else: 244 return 'i386' 245 246 247def get_current_kernel_arch(): 248 """Get the machine architecture, now just a wrap of 'uname -m'.""" 249 return os.popen('uname -m').read().rstrip() 250 251 252def get_file_arch(filename): 253 # -L means follow symlinks 254 file_data = system_output('file -L ' + filename) 255 if file_data.count('80386'): 256 return 'i386' 257 return None 258 259 260def kernelexpand(kernel, args=None): 261 # if not (kernel.startswith('http://') or kernel.startswith('ftp://') or os.path.isfile(kernel)): 262 if kernel.find('/') < 0: # contains no path. 263 autodir = os.environ['AUTODIR'] 264 kernelexpand = os.path.join(autodir, 'tools/kernelexpand') 265 if args: 266 kernelexpand += ' ' + args 267 w, r = os.popen2(kernelexpand + ' ' + kernel) 268 269 kernel = r.readline().strip() 270 r.close() 271 w.close() 272 return kernel.split() 273 274 275def count_cpus(): 276 """number of CPUs in the local machine according to /proc/cpuinfo""" 277 f = file('/proc/cpuinfo', 'r') 278 cpus = 0 279 for line in f.readlines(): 280 if line.startswith('processor'): 281 cpus += 1 282 return cpus 283 284 285# Returns total memory in kb 286def memtotal(): 287 memtotal = system_output('grep MemTotal /proc/meminfo') 288 return int(re.search(r'\d+', memtotal).group(0)) 289 290 291def system(cmd, ignorestatus = 0): 292 """os.system replacement 293 294 We have our own definition of system here, as the stock os.system doesn't 295 correctly handle sigpipe 296 (ie things like "yes | head" will hang because yes doesn't get the SIGPIPE). 297 298 Also the stock os.system didn't raise errors based on exit status, this 299 version does unless you explicitly specify ignorestatus=1 300 """ 301 signal.signal(signal.SIGPIPE, signal.SIG_DFL) 302 try: 303 status = os.system(cmd) 304 finally: 305 signal.signal(signal.SIGPIPE, signal.SIG_IGN) 306 307 if ((status != 0) and not ignorestatus): 308 raise CmdError(cmd, status) 309 return status 310 311 312def system_output(command, ignorestatus = 0): 313 """Run command and return its output 314 315 ignorestatus 316 whether to raise a CmdError if command has a nonzero exit status 317 """ 318 (result, data) = commands.getstatusoutput(command) 319 if ((result != 0) and not ignorestatus): 320 raise CmdError, 'command failed: ' + command 321 return data 322 323 324def where_art_thy_filehandles(): 325 """Dump the current list of filehandles""" 326 os.system("ls -l /proc/%d/fd >> /dev/tty" % os.getpid()) 327 328 329def print_to_tty(string): 330 """Output string straight to the tty""" 331 os.system("echo " + string + " >> /dev/tty") 332 333 334def dump_object(object): 335 """Dump an object's attributes and methods 336 337 kind of like dir() 338 """ 339 for item in object.__dict__.iteritems(): 340 print item 341 try: 342 (key,value) = item 343 dump_object(value) 344 except: 345 continue 346 347 348def environ(env_key): 349 """return the requested environment variable, or '' if unset""" 350 if (os.environ.has_key(env_key)): 351 return os.environ[env_key] 352 else: 353 return '' 354 355 356def prepend_path(newpath, oldpath): 357 """prepend newpath to oldpath""" 358 if (oldpath): 359 return newpath + ':' + oldpath 360 else: 361 return newpath 362 363 364def append_path(oldpath, newpath): 365 """append newpath to oldpath""" 366 if (oldpath): 367 return oldpath + ':' + newpath 368 else: 369 return newpath 370 371 372def avgtime_print(dir): 373 """ Calculate some benchmarking statistics. 374 Input is a directory containing a file called 'time'. 375 File contains one-per-line results of /usr/bin/time. 376 Output is average Elapsed, User, and System time in seconds, 377 and average CPU percentage. 378 """ 379 f = open(dir + "/time") 380 user = system = elapsed = cpu = count = 0 381 r = re.compile('([\d\.]*)user ([\d\.]*)system (\d*):([\d\.]*)elapsed (\d*)%CPU') 382 for line in f.readlines(): 383 try: 384 s = r.match(line); 385 user += float(s.group(1)) 386 system += float(s.group(2)) 387 elapsed += (float(s.group(3)) * 60) + float(s.group(4)) 388 cpu += float(s.group(5)) 389 count += 1 390 except: 391 raise ValueError("badly formatted times") 392 393 f.close() 394 return "Elapsed: %0.2fs User: %0.2fs System: %0.2fs CPU: %0.0f%%" % \ 395 (elapsed/count, user/count, system/count, cpu/count) 396 397 398def running_config(): 399 """ 400 Return path of config file of the currently running kernel 401 """ 402 for config in ('/proc/config.gz', \ 403 '/boot/config-%s' % system_output('uname -r') ): 404 if os.path.isfile(config): 405 return config 406 return None 407 408 409def cpu_online_map(): 410 """ 411 Check out the available cpu online map 412 """ 413 cpus = [] 414 for line in open('/proc/cpuinfo', 'r').readlines(): 415 if line.startswith('processor'): 416 cpus.append(line.split()[2]) # grab cpu number 417 return cpus 418 419 420def check_glibc_ver(ver): 421 glibc_ver = commands.getoutput('ldd --version').splitlines()[0] 422 glibc_ver = re.search(r'(\d+\.\d+(\.\d+)?)', glibc_ver).group() 423 if glibc_ver.split('.') < ver.split('.'): 424 raise "Glibc is too old (%s). Glibc >= %s is needed." % \ 425 (glibc_ver, ver) 426 427def read_one_line(filename): 428 return open(filename, 'r').readline().strip() 429 430 431def write_one_line(filename, str): 432 str.rstrip() 433 open(filename, 'w').write(str.rstrip() + "\n") 434 435 436def human_format(number): 437 # Convert number to kilo / mega / giga format. 438 if number < 1024: 439 return "%d" % number 440 kilo = float(number) / 1024.0 441 if kilo < 1024: 442 return "%.2fk" % kilo 443 meg = kilo / 1024.0 444 if meg < 1024: 445 return "%.2fM" % meg 446 gig = meg / 1024.0 447 return "%.2fG" % gig 448 449