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(
 125                     filename, 'unknown'
 126                 )
 127
 128             path = self.tree_node_by_module.get(loading_module, [])
 129             path.extend([loaded_module])
 130             self.tree.insert(path)
 131             self.tree_node_by_module[loading_module] = path
 132
 133             msg = f'*** Import {loaded_module} from {filename}:{s[x].lineno} in {loading_module}::{loading_function}'
 134             logger.debug(msg)
 135             print(msg)
 136             return
 137         msg = f'*** Import {loaded_module} from ?????'
 138         logger.debug(msg)
 139         print(msg)
 140
 141     def find_importer(self, module: str):
 142         if module in self.tree_node_by_module:
 143             node = self.tree_node_by_module[module]
 144             return node
 145         return []
 146
 147
 148 # # TODO: test this with python 3.8+
 149 # def audit_import_events(event, args):
 150 #     if event == 'import':
 151 #         module = args[0]
 152 #         filename = args[1]
 153 #         sys_path = args[2]
 154 #         sys_meta_path = args[3]
 155 #         sys_path_hooks = args[4]
 156 #         logger.debug(msg)
 157 #         print(msg)
 158
 159
 160 # Audit import events?  Note: this runs early in the lifetime of the
 161 # process (assuming that import bootstrap happens early); config has
 162 # (probably) not yet been loaded or parsed the commandline.  Also,
 163 # some things have probably already been imported while we weren't
 164 # watching so this information may be incomplete.
 165 #
 166 # Also note: move bootstrap up in the global import list to catch
 167 # more import events and have a more complete record.
 168 import_interceptor = None
 169 for arg in sys.argv:
 170     if arg == '--audit_import_events':
 171         import_interceptor = ImportInterceptor()
 172         sys.meta_path = [import_interceptor] + sys.meta_path
 173         # if not hasattr(sys, 'frozen'):
 174         #     if (
 175         #             sys.version_info[0] == 3
 176         #             and sys.version_info[1] >= 8
 177         #     ):
 178         #         sys.addaudithook(audit_import_events)
 179
 180
 181 def dump_all_objects() -> None:
 182     global import_interceptor
 183     messages = {}
 184     all_modules = sys.modules
 185     for obj in object.__subclasses__():
 186         if not hasattr(obj, '__name__'):
 187             continue
 188         klass = obj.__name__
 189         if not hasattr(obj, '__module__'):
 190             continue
 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
 196             else:
 197                 mod_name = mod.__name__
 198             if hasattr(mod, '__file__'):
 199                 mod_file = mod.__file__
 200             else:
 201                 mod_file = 'unknown'
 202             if import_interceptor is not None:
 203                 import_path = import_interceptor.find_importer(mod_name)
 204             else:
 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])
 212         print(messages[x])
 213
 214
 215 def initialize(entry_point):
 216     """
 217     Remember to initialize config, initialize logging, set/log a random
 218     seed, etc... before running main.
 219
 220     """
 221
 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
 227
 228         # Try to figure out the name of the program entry point.  Then
 229         # parse configuration (based on cmdline flags, environment vars
 230         # etc...)
 231         if (
 232             '__globals__' in entry_point.__dict__
 233             and '__file__' in entry_point.__globals__
 234         ):
 235             config.parse(entry_point.__globals__['__file__'])
 236         else:
 237             config.parse(None)
 238
 239         # Initialize logging... and log some remembered messages from
 240         # config module.
 241         logging_utils.initialize_logging(logging.getLogger())
 242         config.late_logging()
 243
 244         # Allow programs that don't bother to override the random seed
 245         # to be replayed via the commandline.
 246         import random
 247
 248         random_seed = config.config['set_random_seed']
 249         if random_seed is not None:
 250             random_seed = random_seed[0]
 251         else:
 252             random_seed = int.from_bytes(os.urandom(4), 'little')
 253
 254         if config.config['show_random_seed']:
 255             msg = f'Global random seed is: {random_seed}'
 256             logger.debug(msg)
 257             print(msg)
 258         random.seed(random_seed)
 259
 260         # Do it, invoke the user's code.  Pay attention to how long it takes.
 261         logger.debug(f'Starting {entry_point.__name__} (program entry point)')
 262         ret = None
 263         import stopwatch
 264
 265         with stopwatch.Timer() as t:
 266             ret = entry_point(*args, **kwargs)
 267         logger.debug(
 268             f'{entry_point.__name__} (program entry point) returned {ret}.'
 269         )
 270
 271         if config.config['dump_all_objects']:
 272             dump_all_objects()
 273
 274         if config.config['audit_import_events']:
 275             global import_interceptor
 276             if import_interceptor is not None:
 277                 print(import_interceptor.tree)
 278
 279         walltime = t()
 280         (utime, stime, cutime, cstime, elapsed_time) = os.times()
 281         logger.debug(
 282             '\n'
 283             f'user: {utime}s\n'
 284             f'system: {stime}s\n'
 285             f'child user: {cutime}s\n'
 286             f'child system: {cstime}s\n'
 287             f'machine uptime: {elapsed_time}s\n'
 288             f'walltime: {walltime}s'
 289         )
 290
 291         # If it doesn't return cleanly, call attention to the return value.
 292         if ret is not None and ret != 0:
 293             logger.error(f'Exit {ret}')
 294         else:
 295             logger.debug(f'Exit {ret}')
 296         sys.exit(ret)
 297
 298     return initialize_wrapper