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