exec_wrapper.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. #!/usr/bin/env python3
  2. # Import as few modules and define as few names as possible at module level before exec'ing the
  3. # file, as all names will be seen by the exec'ed code. To prevent name collisions, all names
  4. # defined or imported here at module level should begin with 'exec_wrapper_'
  5. def exec_wrapper_get_colors():
  6. from collections import namedtuple
  7. return namedtuple('colors',['red','green','yellow','blue','purple'])(*[
  8. (lambda s:s) if exec_wrapper_os.getenv('MMGEN_DISABLE_COLOR') else
  9. (lambda s,n=n:f'\033[{n};1m{s}\033[0m' )
  10. for n in (31,32,33,34,35) ])
  11. def exec_wrapper_init():
  12. if exec_wrapper_os.path.dirname(exec_wrapper_sys.argv[1]) == 'test': # scripts in ./test do overlay setup themselves
  13. exec_wrapper_sys.path[0] = 'test'
  14. else:
  15. exec_wrapper_sys.path.pop(0)
  16. if 'TMUX' in exec_wrapper_os.environ:
  17. del exec_wrapper_os.environ['TMUX']
  18. if exec_wrapper_os.getenv('EXEC_WRAPPER_TRACEBACK'):
  19. try:
  20. exec_wrapper_os.unlink('test.py.err')
  21. except:
  22. pass
  23. exec_wrapper_os.environ['MMGEN_EXEC_WRAPPER'] = '1'
  24. def exec_wrapper_write_traceback(e,exit_val):
  25. import sys,os
  26. exc_line = (
  27. repr(e) if type(e).__name__ in ('MMGenError','MMGenSystemExit') else
  28. '{}: {}'.format( type(e).__name__, e ))
  29. c = exec_wrapper_get_colors()
  30. if os.getenv('EXEC_WRAPPER_TRACEBACK'):
  31. import traceback
  32. cwd = os.getcwd()
  33. sys.path.insert(0,cwd)
  34. from test.overlay import get_overlay_tree_dir
  35. overlay_path_pfx = os.path.relpath(get_overlay_tree_dir(cwd)) + '/'
  36. def fixup_fn(fn_in):
  37. from mmgen.util2 import removeprefix,removesuffix
  38. fn = removeprefix(removeprefix(fn_in,cwd+'/'),overlay_path_pfx)
  39. return removesuffix(fn,'_orig.py') + '.py' if fn.endswith('_orig.py') else fn
  40. # Python 3.9:
  41. # fn = fn_in.removeprefix(cwd+'/').removeprefix('test/overlay/tree/')
  42. # return fn.removesuffix('_orig.py') + '.py' if fn.endswith('_orig.py') else fn
  43. def gen_output():
  44. yield 'Traceback (most recent call last):'
  45. for e in traceback.extract_tb(sys.exc_info()[2]):
  46. yield ' File "{f}", line {l}, in {n}\n {L}'.format(
  47. f = exec_wrapper_execed_file if e.filename == '<string>' else fixup_fn(e.filename),
  48. l = '(scrubbed)' if os.getenv('MMGEN_TEST_SUITE_DETERMINISTIC') else e.lineno,
  49. n = e.name,
  50. L = e.line or 'N/A' )
  51. tb_lines = list( gen_output() )
  52. if 'SystemExit' in exc_line:
  53. tb_lines.pop()
  54. sys.stdout.write('{}\n{}\n'.format( c.yellow( '\n'.join(tb_lines) ), c.red(exc_line) ))
  55. from test.include.common import test_py_error_fn
  56. with open(test_py_error_fn,'w') as fp:
  57. fp.write('\n'.join(tb_lines + [exc_line]))
  58. print(c.blue('{} script exited with error').format(
  59. 'Test' if os.path.dirname(sys.argv[0]) == 'test' else 'Spawned' ))
  60. else:
  61. sys.stdout.write( c.purple((f'NONZERO_EXIT[{exit_val}]: ' if exit_val else '') + exc_line) + '\n' )
  62. def exec_wrapper_end_msg():
  63. if (
  64. exec_wrapper_os.getenv('EXEC_WRAPPER_SPAWN')
  65. and not exec_wrapper_os.getenv('MMGEN_TEST_SUITE_DETERMINISTIC') ):
  66. c = exec_wrapper_get_colors()
  67. # write to stdout to ensure script output gets to terminal first
  68. exec_wrapper_sys.stdout.write(c.blue('Runtime: {:0.5f} secs\n'.format(
  69. exec_wrapper_time.time() - exec_wrapper_tstart )))
  70. def exec_wrapper_tracemalloc_setup():
  71. if exec_wrapper_os.getenv('MMGEN_TRACEMALLOC'):
  72. exec_wrapper_os.environ['PYTHONTRACEMALLOC'] = '1'
  73. import tracemalloc
  74. tracemalloc.start()
  75. exec_wrapper_sys.stderr.write("INFO → Appending memory allocation stats to 'tracemalloc.log'\n")
  76. def exec_wrapper_tracemalloc_log():
  77. if exec_wrapper_os.getenv('MMGEN_TRACEMALLOC'):
  78. import tracemalloc,re
  79. snapshot = tracemalloc.take_snapshot()
  80. stats = snapshot.statistics('lineno')
  81. depth = 100
  82. col1w = 100
  83. with open('tracemalloc.log','a') as fp:
  84. fp.write('##### TOP {} {} #####\n'.format(depth,' '.join(exec_wrapper_sys.argv)))
  85. for stat in stats[:depth]:
  86. frame = stat.traceback[0]
  87. fn = re.sub(r'.*\/site-packages\/|.*\/mmgen\/test\/overlay\/tree\/','',frame.filename)
  88. fn = re.sub(r'.*\/mmgen\/test\/','test/',fn)
  89. fp.write('{f:{w}} {s:>8.2f} KiB\n'.format(
  90. f = f'{fn}:{frame.lineno}:',
  91. s = stat.size/1024,
  92. w = col1w ))
  93. fp.write('{f:{w}} {s:8.2f} KiB\n\n'.format(
  94. f = 'TOTAL {}:'.format(' '.join(exec_wrapper_sys.argv))[:col1w],
  95. s = sum(stat.size for stat in stats) / 1024,
  96. w = col1w ))
  97. import sys as exec_wrapper_sys
  98. import os as exec_wrapper_os
  99. import time as exec_wrapper_time
  100. exec_wrapper_init() # sets sys.path[0], runs overlay_setup()
  101. exec_wrapper_tracemalloc_setup()
  102. # import mmgen mods only after overlay setup!
  103. from mmgen.devinit import init_dev as exec_wrapper_init_dev
  104. exec_wrapper_init_dev()
  105. exec_wrapper_tstart = exec_wrapper_time.time()
  106. try:
  107. exec_wrapper_sys.argv.pop(0)
  108. exec_wrapper_execed_file = exec_wrapper_sys.argv[0]
  109. with open(exec_wrapper_execed_file) as fp:
  110. exec(fp.read())
  111. except SystemExit as e:
  112. if e.code != 0:
  113. exec_wrapper_write_traceback(e,e.code)
  114. else:
  115. exec_wrapper_tracemalloc_log()
  116. exec_wrapper_end_msg()
  117. exec_wrapper_sys.exit(e.code)
  118. except Exception as e:
  119. exit_val = e.mmcode if hasattr(e,'mmcode') else e.code if hasattr(e,'code') else 1
  120. exec_wrapper_write_traceback(e,exit_val)
  121. exec_wrapper_sys.exit(exit_val)
  122. exec_wrapper_tracemalloc_log()
  123. exec_wrapper_end_msg()