bootstrap.py

   1 #!/usr/bin/env python3
   2
   3 import functools
   4 import logging
   5 import os
   6 from inspect import stack
   7 import sys
   8
   9 # This module is commonly used by others in here and should avoid
  10 # taking any unnecessary dependencies back on them.
  11
  12 from argparse_utils import ActionNoYes
  13 import config
  14 import logging_utils
  15
  16 logger = logging.getLogger(__name__)
  17
  18 args = config.add_commandline_args(
  19     f'Bootstrap ({__file__})',
  20     'Args related to python program bootstrapper and Swiss army knife')
  21 args.add_argument(
  22     '--debug_unhandled_exceptions',
  23     action=ActionNoYes,
  24     default=False,
  25     help='Break into pdb on top level unhandled exceptions.'
  26 )
  27 args.add_argument(
  28     '--show_random_seed',
  29     action=ActionNoYes,
  30     default=False,
  31     help='Should we display (and log.debug) the global random seed?'
  32 )
  33 args.add_argument(
  34     '--set_random_seed',
  35     type=int,
  36     nargs=1,
  37     default=None,
  38     metavar='SEED_INT',
  39     help='Override the global random seed with a particular number.'
  40 )
  41 args.add_argument(
  42     '--dump_all_objects',
  43     action=ActionNoYes,
  44     default=False,
  45     help='Should we dump the Python import tree before main?'
  46 )
  47 args.add_argument(
  48     '--audit_import_events',
  49     action=ActionNoYes,
  50     default=False,
  51     help='Should we audit all import events?',
  52 )
  53
  54
  55 original_hook = sys.excepthook
  56
  57
  58 def handle_uncaught_exception(exc_type, exc_value, exc_tb):
  59     """
  60     Top-level exception handler for exceptions that make it past any exception
  61     handlers in the python code being run.  Logs the error and stacktrace then
  62     maybe attaches a debugger.
  63
  64     """
  65     global original_hook
  66     msg = f'Unhandled top level exception {exc_type}'
  67     logger.exception(msg)
  68     print(msg, file=sys.stderr)
  69     if issubclass(exc_type, KeyboardInterrupt):
  70         sys.__excepthook__(exc_type, exc_value, exc_tb)
  71         return
  72     else:
  73         if (
  74                 not sys.stderr.isatty() or
  75                 not sys.stdin.isatty()
  76         ):
  77             # stdin or stderr is redirected, just do the normal thing
  78             original_hook(exc_type, exc_value, exc_tb)
  79         else:
  80             # a terminal is attached and stderr is not redirected, maybe debug.
  81             import traceback
  82             traceback.print_exception(exc_type, exc_value, exc_tb)
  83             if config.config['debug_unhandled_exceptions']:
  84                 import pdb
  85                 logger.info("Invoking the debugger...")
  86                 pdb.pm()
  87             else:
  88                 original_hook(exc_type, exc_value, exc_tb)
  89
  90
  91 class ImportInterceptor(object):
  92     def __init__(self):
  93         import collect.trie
  94         self.module_by_filename_cache = {}
  95         self.repopulate_modules_by_filename()
  96         self.tree = collect.trie.Trie()
  97         self.tree_node_by_module = {}
  98
  99     def repopulate_modules_by_filename(self):
 100         self.module_by_filename_cache.clear()
 101         for mod in sys.modules:
 102             if hasattr(sys.modules[mod], '__file__'):
 103                 fname = getattr(sys.modules[mod], '__file__')
 104             else:
 105                 fname = 'unknown'
 106             self.module_by_filename_cache[fname] = mod
 107
 108     def should_ignore_filename(self, filename: str) -> bool:
 109         return 'importlib' in filename or 'six.py' in filename
 110
 111     def find_spec(self, loaded_module, path=None, target=None):
 112         s = stack()
 113         for x in range(3, len(s)):
 114             filename = s[x].filename
 115             if self.should_ignore_filename(filename):
 116                 continue
 117
 118             loading_function = s[x].function
 119             if filename in self.module_by_filename_cache:
 120                 loading_module = self.module_by_filename_cache[filename]
 121             else:
 122                 self.repopulate_modules_by_filename()
 123                 loading_module = self.module_by_filename_cache.get(filename, 'unknown')
 124
 125             path = self.tree_node_by_module.get(loading_module, [])
 126             path.extend([loaded_module])
 127             self.tree.insert(path)
 128             self.tree_node_by_module[loading_module] = path
 129
 130             msg = f'*** Import {loaded_module} from {filename}:{s[x].lineno} in {loading_module}::{loading_function}'
 131             logger.debug(msg)
 132             print(msg)
 133             return
 134         msg = f'*** Import {loaded_module} from ?????'
 135         logger.debug(msg)
 136         print(msg)
 137
 138     def find_importer(self, module: str):
 139         if module in self.tree_node_by_module:
 140             node = self.tree_node_by_module[module]
 141             return node
 142         return []
 143
 144
 145 # # TODO: test this with python 3.8+
 146 # def audit_import_events(event, args):
 147 #     if event == 'import':
 148 #         module = args[0]
 149 #         filename = args[1]
 150 #         sys_path = args[2]
 151 #         sys_meta_path = args[3]
 152 #         sys_path_hooks = args[4]
 153 #         logger.debug(msg)
 154 #         print(msg)
 155
 156
 157 # Audit import events?  Note: this runs early in the lifetime of the
 158 # process (assuming that import bootstrap happens early); config has
 159 # (probably) not yet been loaded or parsed the commandline.  Also,
 160 # some things have probably already been imported while we weren't
 161 # watching so this information may be incomplete.
 162 #
 163 # Also note: move bootstrap up in the global import list to catch
 164 # more import events and have a more complete record.
 165 import_interceptor = None
 166 for arg in sys.argv:
 167     if arg == '--audit_import_events':
 168         import_interceptor = ImportInterceptor()
 169         sys.meta_path = [import_interceptor] + sys.meta_path
 170         # if not hasattr(sys, 'frozen'):
 171         #     if (
 172         #             sys.version_info[0] == 3
 173         #             and sys.version_info[1] >= 8
 174         #     ):
 175         #         sys.addaudithook(audit_import_events)
 176
 177
 178 def dump_all_objects() -> None:
 179     global import_interceptor
 180     messages = {}
 181     all_modules = sys.modules
 182     for obj in object.__subclasses__():
 183         if not hasattr(obj, '__name__'):
 184             continue
 185         klass = obj.__name__
 186         if not hasattr(obj, '__module__'):
 187             continue
 188         class_mod_name = obj.__module__
 189         if class_mod_name in all_modules:
 190             mod = all_modules[class_mod_name]
 191             if not hasattr(mod, '__name__'):
 192                 mod_name = class_mod_name
 193             else:
 194                 mod_name = mod.__name__
 195             if hasattr(mod, '__file__'):
 196                 mod_file = mod.__file__
 197             else:
 198                 mod_file = 'unknown'
 199             if import_interceptor is not None:
 200                 import_path = import_interceptor.find_importer(mod_name)
 201             else:
 202                 import_path = 'unknown'
 203             msg = f'{class_mod_name}::{klass} ({mod_file})'
 204             if import_path != 'unknown' and len(import_path) > 0:
 205                 msg += f' imported by {import_path}'
 206             messages[f'{class_mod_name}::{klass}'] = msg
 207     for x in sorted(messages.keys()):
 208         logger.debug(messages[x])
 209         print(messages[x])
 210
 211
 212 def initialize(entry_point):
 213     """
 214     Remember to initialize config, initialize logging, set/log a random
 215     seed, etc... before running main.
 216
 217     """
 218     @functools.wraps(entry_point)
 219     def initialize_wrapper(*args, **kwargs):
 220         # Hook top level unhandled exceptions, maybe invoke debugger.
 221         if sys.excepthook == sys.__excepthook__:
 222             sys.excepthook = handle_uncaught_exception
 223
 224         # Try to figure out the name of the program entry point.  Then
 225         # parse configuration (based on cmdline flags, environment vars
 226         # etc...)
 227         if (
 228                 '__globals__' in entry_point.__dict__ and
 229                 '__file__' in entry_point.__globals__
 230         ):
 231             config.parse(entry_point.__globals__['__file__'])
 232         else:
 233             config.parse(None)
 234
 235         # Initialize logging... and log some remembered messages from
 236         # config module.
 237         logging_utils.initialize_logging(logging.getLogger())
 238         config.late_logging()
 239
 240         # Allow programs that don't bother to override the random seed
 241         # to be replayed via the commandline.
 242         import random
 243         random_seed = config.config['set_random_seed']
 244         if random_seed is not None:
 245             random_seed = random_seed[0]
 246         else:
 247             random_seed = int.from_bytes(os.urandom(4), 'little')
 248
 249         if config.config['show_random_seed']:
 250             msg = f'Global random seed is: {random_seed}'
 251             logger.debug(msg)
 252             print(msg)
 253         random.seed(random_seed)
 254
 255         # Do it, invoke the user's code.  Pay attention to how long it takes.
 256         logger.debug(f'Starting {entry_point.__name__} (program entry point)')
 257         ret = None
 258         import stopwatch
 259         with stopwatch.Timer() as t:
 260             ret = entry_point(*args, **kwargs)
 261         logger.debug(
 262             f'{entry_point.__name__} (program entry point) returned {ret}.'
 263         )
 264
 265         if config.config['dump_all_objects']:
 266             dump_all_objects()
 267
 268         if config.config['audit_import_events']:
 269             global import_interceptor
 270             if import_interceptor is not None:
 271                 print(import_interceptor.tree)
 272
 273         walltime = t()
 274         (utime, stime, cutime, cstime, elapsed_time) = os.times()
 275         logger.debug('\n'
 276                      f'user: {utime}s\n'
 277                      f'system: {stime}s\n'
 278                      f'child user: {cutime}s\n'
 279                      f'child system: {cstime}s\n'
 280                      f'machine uptime: {elapsed_time}s\n'
 281                      f'walltime: {walltime}s')
 282
 283         # If it doesn't return cleanly, call attention to the return value.
 284         if ret is not None and ret != 0:
 285             logger.error(f'Exit {ret}')
 286         else:
 287             logger.debug(f'Exit {ret}')
 288         sys.exit(ret)
 289     return initialize_wrapper