8 from inspect import stack
9 from typing import List
13 from argparse_utils import ActionNoYes
15 # This module is commonly used by others in here and should avoid
16 # taking any unnecessary dependencies back on them.
19 logger = logging.getLogger(__name__)
21 args = config.add_commandline_args(
22 f'Bootstrap ({__file__})',
23 'Args related to python program bootstrapper and Swiss army knife',
26 '--debug_unhandled_exceptions',
29 help='Break into pdb on top level unhandled exceptions.',
35 help='Should we display (and log.debug) the global random seed?',
43 help='Override the global random seed with a particular number.',
49 help='Should we dump the Python import tree before main?',
52 '--audit_import_events',
55 help='Should we audit all import events?',
61 help='Should we run cProfile on this code?',
67 help='Should we record/report on memory utilization?',
70 original_hook = sys.excepthook
73 def handle_uncaught_exception(exc_type, exc_value, exc_tb):
75 Top-level exception handler for exceptions that make it past any exception
76 handlers in the python code being run. Logs the error and stacktrace then
77 maybe attaches a debugger.
81 msg = f'Unhandled top level exception {exc_type}'
83 print(msg, file=sys.stderr)
84 if issubclass(exc_type, KeyboardInterrupt):
85 sys.__excepthook__(exc_type, exc_value, exc_tb)
88 if not sys.stderr.isatty() or not sys.stdin.isatty():
89 # stdin or stderr is redirected, just do the normal thing
90 original_hook(exc_type, exc_value, exc_tb)
92 # a terminal is attached and stderr is not redirected, maybe debug.
95 traceback.print_exception(exc_type, exc_value, exc_tb)
96 if config.config['debug_unhandled_exceptions']:
99 logger.info("Invoking the debugger...")
102 original_hook(exc_type, exc_value, exc_tb)
105 class ImportInterceptor(importlib.abc.MetaPathFinder):
109 self.module_by_filename_cache = {}
110 self.repopulate_modules_by_filename()
111 self.tree = collect.trie.Trie()
112 self.tree_node_by_module = {}
114 def repopulate_modules_by_filename(self):
115 self.module_by_filename_cache.clear()
116 for mod in sys.modules:
117 if hasattr(sys.modules[mod], '__file__'):
118 fname = getattr(sys.modules[mod], '__file__')
121 self.module_by_filename_cache[fname] = mod
123 def should_ignore_filename(self, filename: str) -> bool:
124 return 'importlib' in filename or 'six.py' in filename
126 def find_module(self, fullname, path):
127 raise Exception("This method has been deprecated since Python 3.4, please upgrade.")
129 def find_spec(self, loaded_module, path=None, target=None):
131 for x in range(3, len(s)):
132 filename = s[x].filename
133 if self.should_ignore_filename(filename):
136 loading_function = s[x].function
137 if filename in self.module_by_filename_cache:
138 loading_module = self.module_by_filename_cache[filename]
140 self.repopulate_modules_by_filename()
141 loading_module = self.module_by_filename_cache.get(filename, 'unknown')
143 path = self.tree_node_by_module.get(loading_module, [])
144 path.extend([loaded_module])
145 self.tree.insert(path)
146 self.tree_node_by_module[loading_module] = path
148 msg = f'*** Import {loaded_module} from {filename}:{s[x].lineno} in {loading_module}::{loading_function}'
152 msg = f'*** Import {loaded_module} from ?????'
156 def invalidate_caches(self):
159 def find_importer(self, module: str):
160 if module in self.tree_node_by_module:
161 node = self.tree_node_by_module[module]
166 # Audit import events? Note: this runs early in the lifetime of the
167 # process (assuming that import bootstrap happens early); config has
168 # (probably) not yet been loaded or parsed the commandline. Also,
169 # some things have probably already been imported while we weren't
170 # watching so this information may be incomplete.
172 # Also note: move bootstrap up in the global import list to catch
173 # more import events and have a more complete record.
174 import_interceptor = None
176 if arg == '--audit_import_events':
177 import_interceptor = ImportInterceptor()
178 sys.meta_path.insert(0, import_interceptor)
181 def dump_all_objects() -> None:
182 global import_interceptor
184 all_modules = sys.modules
185 for obj in object.__subclasses__():
186 if not hasattr(obj, '__name__'):
189 if not hasattr(obj, '__module__'):
191 class_mod_name = obj.__module__
192 if class_mod_name in all_modules:
193 mod = all_modules[class_mod_name]
194 if not hasattr(mod, '__name__'):
195 mod_name = class_mod_name
197 mod_name = mod.__name__
198 if hasattr(mod, '__file__'):
199 mod_file = mod.__file__
202 if import_interceptor is not None:
203 import_path = import_interceptor.find_importer(mod_name)
205 import_path = 'unknown'
206 msg = f'{class_mod_name}::{klass} ({mod_file})'
207 if import_path != 'unknown' and len(import_path) > 0:
208 msg += f' imported by {import_path}'
209 messages[f'{class_mod_name}::{klass}'] = msg
210 for x in sorted(messages.keys()):
211 logger.debug(messages[x])
215 def initialize(entry_point):
217 Remember to initialize config, initialize logging, set/log a random
218 seed, etc... before running main.
222 @functools.wraps(entry_point)
223 def initialize_wrapper(*args, **kwargs):
224 # Hook top level unhandled exceptions, maybe invoke debugger.
225 if sys.excepthook == sys.__excepthook__:
226 sys.excepthook = handle_uncaught_exception
228 # Try to figure out the name of the program entry point. Then
229 # parse configuration (based on cmdline flags, environment vars
231 if '__globals__' in entry_point.__dict__ and '__file__' in entry_point.__globals__:
232 config.parse(entry_point.__globals__['__file__'])
236 if config.config['trace_memory']:
241 # Initialize logging... and log some remembered messages from
243 logging_utils.initialize_logging(logging.getLogger())
244 config.late_logging()
246 # Maybe log some info about the python interpreter itself.
248 f'Platform: {sys.platform}, maxint=0x{sys.maxsize:x}, byteorder={sys.byteorder}'
250 logger.debug(f'Python interpreter version: {sys.version}')
251 logger.debug(f'Python implementation: {sys.implementation}')
252 logger.debug(f'Python C API version: {sys.api_version}')
253 logger.debug(f'Python path: {sys.path}')
255 # Allow programs that don't bother to override the random seed
256 # to be replayed via the commandline.
259 random_seed = config.config['set_random_seed']
260 if random_seed is not None:
261 random_seed = random_seed[0]
263 random_seed = int.from_bytes(os.urandom(4), 'little')
265 if config.config['show_random_seed']:
266 msg = f'Global random seed is: {random_seed}'
269 random.seed(random_seed)
271 # Do it, invoke the user's code. Pay attention to how long it takes.
272 logger.debug(f'Starting {entry_point.__name__} (program entry point)')
276 if config.config['run_profiler']:
278 from pstats import SortKey
280 with stopwatch.Timer() as t:
282 "ret = entry_point(*args, **kwargs)",
289 with stopwatch.Timer() as t:
290 ret = entry_point(*args, **kwargs)
292 logger.debug(f'{entry_point.__name__} (program entry point) returned {ret}.')
294 if config.config['trace_memory']:
295 snapshot = tracemalloc.take_snapshot()
296 top_stats = snapshot.statistics('lineno')
298 print("--trace_memory's top 10 memory using files:")
299 for stat in top_stats[:10]:
302 if config.config['dump_all_objects']:
305 if config.config['audit_import_events']:
306 global import_interceptor
307 if import_interceptor is not None:
308 print(import_interceptor.tree)
311 (utime, stime, cutime, cstime, elapsed_time) = os.times()
315 f'system: {stime}s\n'
316 f'child user: {cutime}s\n'
317 f'child system: {cstime}s\n'
318 f'machine uptime: {elapsed_time}s\n'
319 f'walltime: {walltime}s'
322 # If it doesn't return cleanly, call attention to the return value.
323 if ret is not None and ret != 0:
324 logger.error(f'Exit {ret}')
326 logger.debug(f'Exit {ret}')
329 return initialize_wrapper