setup_util.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. import re
  2. import os
  3. import subprocess
  4. # Replaces all text found using the regular expression to_replace with the supplied replacement.
  5. def replace_text(file, to_replace, replacement):
  6. with open(file, "r") as conf:
  7. contents = conf.read()
  8. replaced_text = re.sub(to_replace, replacement, contents)
  9. with open(file, "w") as f:
  10. f.write(replaced_text)
  11. # Replaces the current process environment with the one found in
  12. # config file. Retains a few original vars (HOME,PATH, etc) by default.
  13. # Optionally allows specification of a command to be run before loading
  14. # the environment, to allow the framework to set environment variables
  15. # Note: This command *cannot* print to stdout!
  16. #
  17. # Note: This will not replace the sudo environment (e.g. subprocess.check_call("sudo <command>")).
  18. # If you must use sudo, consider sudo sh -c ". <config> && your_command"
  19. def replace_environ(config=None, root=None, print_result=False, command='true'):
  20. # Clean up our current environment, preserving some important items
  21. mini_environ = {}
  22. for envname in ['HOME', 'PATH', 'USER', 'LD_LIBRARY_PATH', 'PYTHONPATH', 'FWROOT']:
  23. if envname in os.environ:
  24. mini_environ[envname] = os.environ[envname]
  25. for key in os.environ:
  26. if key.startswith('TFB_'): # Any TFB_ variables are preserved
  27. mini_environ[key] = os.environ[key]
  28. os.environ.clear()
  29. # Use FWROOT if explicitely provided
  30. if root is not None:
  31. mini_environ['FWROOT']=root
  32. # Run command, source config file, and store resulting environment
  33. setup_env = "%s && . %s && env" % (command, config)
  34. env = subprocess.check_output(setup_env, shell=True, env=mini_environ,
  35. executable='/bin/bash')
  36. for line in env.split('\n'):
  37. try:
  38. key, value = line.split('=', 1)
  39. # If we already have this TFB_ variable, do not overwrite
  40. if key.startswith('TFB_') and key in mini_environ:
  41. os.environ[key]=mini_environ[key]
  42. else:
  43. os.environ[key]=value
  44. except:
  45. if not line: # Don't warn for empty line
  46. continue
  47. print "WARN: Line '%s' from '%s' is not an environment variable" % (line, config)
  48. continue
  49. if print_result:
  50. out = subprocess.check_output('env', shell=True, executable='/bin/bash')
  51. print "Environment after loading %s" %config
  52. print out
  53. # Queries the shell for the value of FWROOT
  54. def get_fwroot():
  55. try:
  56. # Use printf to avoid getting a newline
  57. # Redirect to avoid stderr printing
  58. fwroot = subprocess.check_output('printf $FWROOT 2> /dev/null', shell=True, executable='/bin/bash')
  59. return fwroot
  60. except subprocess.CalledProcessError:
  61. return "";
  62. # Turns absolute path into path relative to FWROOT
  63. # Assumes path is underneath FWROOT, not above
  64. #
  65. # Useful for clean presentation of paths
  66. # e.g. /foo/bar/benchmarks/go/bash_profile.sh
  67. # v.s. FWROOT/go/bash_profile.sh
  68. def path_relative_to_root(path):
  69. # Requires bash shell parameter expansion
  70. return subprocess.check_output("D=%s && printf \"${D#%s}\""%(path, get_fwroot()), shell=True, executable='/bin/bash')