AdvancedLogger

advanced_logger.py

import logging


class AdvancedLogger:
    """
    AdvancedLogger can be used for fine grained logging temporary modifications

    How it works:
    Modifications will be enabled when context_manager/decorator starts working and be reverted after

    Usage:
    AdvancedLogger can be used
    - as decorator `@AdvancedLogger()`
    - as context manager `with  AdvancedLogger():`

    It has three main functions/features: 
    - disable loggers and it's handlers by using disable_logger= argument
    - enable/change loggers and it's handlers by using enable_logger= argument
    - disable specific handlers for all loggers, by using  disable_handler= argument
    
    All features they can be used together

    1. Disable loggers

    #Disable all logging handlers
    AdvancedLogger(disable_logger='__all__)

    #Disable specific logger for example "stripe"
    AdvancedLogger(disable_logger="stripe")
    AdvancedLogger(disable_logger=["stripe"])
    AdvancedLogger(disable_logger={"stripe": None})

    #Disable specific logger handler, for example for stripe logger disable console
    AdvancedLogger(disable_logger={"stripe": "console"})
    AdvancedLogger(disable_logger={"stripe": ["console", "console2"]})

    2. Enable/Set loggers

    # Set level for "stripe" logger to 50
    AdvancedLogger(enable_logger={"stripe": 50})
    AdvancedLogger(enable_logger={"stripe": {"level": 50, "propagate": True"}

    # Add already registered handlers (which other loggers have) to your logger
    AdvancedLogger(enable_logger={"stripe": {"handlers": "console"}
    AdvancedLogger(enable_logger={"stripe": {"level": 50, "propagate": True", "handlers": ["console"]}
    
    3. Disable handlers for all loggers
    AdvancedLogger(disable_handler="console")
    AdvancedLogger(disable_handler=["console"])

    Notes:
    Should not be used in production, because
    - not well tested
    - uses private attributes of logging
    - not thread safe

    """
    ACTION_DISABLE_LOGGER_HANDLERS = 'disable_logger_handlers'
    ACTION_DISABLE_LOGGER = 'disable_logger'
    ACTION_DISABLE_ALL = 'disable_all'

    STEP_DISABLE_ALL = ACTION_DISABLE_ALL
    STEP_DISABLE_HANDLERS = ACTION_DISABLE_LOGGER_HANDLERS
    STEP_RESET_PROPAGATE = 'reset_propagate'

    STEP_ADD_HANDLERS = 'add_handlers'
    STEP_SET_LEVEL = 'set_level'
    STEP_SET_PROPAGATE = 'set_level'

    STEP_MUTE_HANDLER = 'mute_handler'

    def __init__(self, disable_logger=None, enable_logger=None, disable_handler=None):
        self.disable_config = disable_logger
        self.disable_steps = self.get_disable_steps(disable_logger)

        self.enable_config = enable_logger
        self.enable_steps = self.get_enable_steps(enable_logger)

        self.disable_handler_config = disable_handler
        self.disable_handler_steps = self.get_disable_handler_steps(disable_handler)

        self.silent_filter = type('', (logging.Filter,), {'filter': staticmethod(lambda r: False)})

    def __call__(self, func):
        """allowed to be used as decorator"""
        def decorator(*args, **kwargs):
            with self:
                return func(*args, **kwargs)
        return decorator

    def process_disable_step(self, step_dict):
        step = step_dict['step']
        if step == self.STEP_DISABLE_ALL:
            logging.disable(logging.CRITICAL)
        elif step == self.STEP_RESET_PROPAGATE:
            logger = step_dict['logger']
            logger.propagate = False
        elif step == self.STEP_DISABLE_HANDLERS:
            logger = step_dict['logger']
            for handler in step_dict['handlers'][:]:
                logger.removeHandler(handler)
        else:
            raise ValueError('Unknown step %s' % step)

    def revert_disable_step(self, step_dict):
        step = step_dict['step']
        if step == self.STEP_DISABLE_ALL:
            logging.disable(logging.NOTSET)
        elif step == self.STEP_RESET_PROPAGATE:
            logger = step_dict['logger']
            logger.propagate = step_dict['propagate']
        elif step == self.STEP_DISABLE_HANDLERS:
            logger = step_dict['logger']
            for handler in step_dict['handlers'][:]:
                logger.addHandler(handler)
        else:
            raise ValueError('Unknown step %s' % step)

    def process_enable_step(self, step_dict):
        step = step_dict['step']
        logger = step_dict['logger']
        if step == self.STEP_SET_LEVEL:
            logger.level = step_dict['level']
        elif step == self.STEP_SET_PROPAGATE:
            logger.propagate = step_dict['propagate']
        elif step == self.STEP_ADD_HANDLERS:
            for handler in step_dict['handlers']:
                logger.addHandler(handler)
        else:
            raise ValueError('Unknown step %s' % step)

    def revert_enable_step(self, step_dict):
        step = step_dict['step']
        logger = step_dict['logger']
        if step == self.STEP_SET_LEVEL:
            logger.level = step_dict['orig_level']
        elif step == self.STEP_SET_PROPAGATE:
            logger.propagate = step_dict['orig_propagate']
        elif step == self.STEP_ADD_HANDLERS:
            for handler in step_dict['handlers']:
                logger.removeHandler(handler)
        else:
            raise ValueError('Unknown step %s' % step)

    def process_disable_handler_step(self, step_dict):
        step = step_dict['step']
        if step == self.STEP_MUTE_HANDLER:
            handlers = step_dict['handlers']
            for handler in handlers:
                handler.addFilter(self.silent_filter)
        else:
            raise ValueError('Unknown step %s' % step)

    def revert_disable_handler_step(self, step_dict):
        step = step_dict['step']
        if step == self.STEP_MUTE_HANDLER:
            handlers = step_dict['handlers']
            for handler in handlers:
                handler.removeFilter(self.silent_filter)
        else:
            raise ValueError('Unknown step %s' % step)

    def __enter__(self):
        for step in self.enable_steps:
            self.process_enable_step(step)
        for step in self.disable_steps:
            self.process_disable_step(step)
        for step in self.disable_handler_steps:
            self.process_disable_handler_step(step)

    def __exit__(self, exit_type, exit_value, exit_traceback):
        for step in self.enable_steps:
            self.revert_enable_step(step)
        for step in self.disable_steps:
            self.revert_disable_step(step)
        for step in self.disable_handler_steps:
            self.revert_disable_handler_step(step)

    def prepare_disable_action_steps(self, action, **kwargs):
        steps = []
        if action == self.ACTION_DISABLE_ALL:
            assert kwargs == {}
            steps.append({
                'step': self.STEP_DISABLE_ALL
            })
        elif action == self.ACTION_DISABLE_LOGGER:
            logger_name = kwargs.pop('logger')
            logger = logging.getLogger(logger_name)
            assert kwargs == {}
            logger_handlers = logger.handlers[:]
            if logger_handlers:
                steps.append({
                    'step': self.STEP_DISABLE_HANDLERS,
                    'logger': logger,
                    'handlers': logger_handlers
                })
            if logger.propagate:
                steps.append({
                    'step': self.STEP_RESET_PROPAGATE,
                    'logger': logger,
                    'propagate': logger.propagate
                })
        elif action == self.ACTION_DISABLE_LOGGER_HANDLERS:
            logger_name = kwargs.pop('logger')
            handler_names = kwargs.pop('handlers')
            logger = logging.getLogger(logger_name)
            assert kwargs == {}
            logger_handlers = [handler for handler in logger.handlers if handler.name in handler_names]
            if logger_handlers:
                steps.append({
                    'step': self.STEP_DISABLE_HANDLERS,
                    'logger': logger,
                    'handlers': logger_handlers,
                })
        return steps

    def get_disable_steps(self, config):
        all_steps = []
        if config is None:
            pass
        elif config is '__all__':
            steps = self.prepare_disable_action_steps(self.ACTION_DISABLE_ALL)
            all_steps.extend(steps)
        elif isinstance(config, str):
            steps = self.prepare_disable_action_steps(self.ACTION_DISABLE_LOGGER, logger=config)
            all_steps.extend(steps)
        elif isinstance(config, (list, tuple)):
            for logger in config:
                steps = self.get_disable_steps(config=logger)
                all_steps.extend(steps)
        elif isinstance(config, dict):
            for logger, handlers in config.items():
                if handlers is None:
                    steps = self.prepare_disable_action_steps(self.ACTION_DISABLE_LOGGER, logger=logger)
                elif isinstance(handlers, str):
                    steps = self.prepare_disable_action_steps(self.ACTION_DISABLE_LOGGER_HANDLERS,
                                                              logger=logger, handlers=[handlers])
                elif isinstance(handlers, (list, tuple)):
                    steps = self.prepare_disable_action_steps(self.ACTION_DISABLE_LOGGER_HANDLERS,
                                                              logger=logger, handlers=handlers)
                else:
                    raise ValueError('Unsupported config format')
                all_steps.extend(steps)
        else:
            raise ValueError('Unsupported config type')

        return all_steps

    def get_enable_steps(self, config):
        all_steps = []
        if config is None:
            pass
        else:
            assert isinstance(config, dict)
            for logger_name, logger_config in config.items():
                logger = logging.getLogger(logger_name)
                if isinstance(logger_config, dict):
                    propagate = logger_config.get('propagate')
                    if propagate is not None and propagate != logger.propagate:
                        all_steps.append({
                            'step': self.STEP_SET_PROPAGATE,
                            'logger': logger,
                            'propagate': propagate,
                            'orig_propagate': logger.propagate
                        })

                    level = logger_config.get('level')
                    if level is not None:
                        steps = self.get_enable_steps(config={logger_name: level})
                        all_steps.extend(steps)

                    handler_names = logger_config.get('handlers')
                    if handler_names:
                        new_handlers = self.get_handlers(handler_names, exclude_handlers=logger.handlers)
                        if new_handlers:
                            all_steps.append({
                                'step': self.STEP_ADD_HANDLERS,
                                'logger': logger,
                                'handlers': new_handlers,
                            })
                else:
                    assert isinstance(logger_config, int) is True, logger_config
                    if logger.level != logger_config:
                        all_steps.append({
                            'step': self.STEP_SET_LEVEL,
                            'logger': logger,
                            'level': logger_config,
                            'orig_level': logger.level
                        })

        return all_steps

    def get_disable_handler_steps(self, config):
        all_steps = []
        if config is None:
            pass
        else:
            assert isinstance(config, (list, tuple, str))
            handlers = self.get_handlers(handler_names=config)
            if handlers:
                all_steps.append({
                    'step': self.STEP_MUTE_HANDLER,
                    'handlers': handlers
                })

        return all_steps

    def get_handlers(self, handler_names, exclude_handlers=None):
        if isinstance(handler_names, str):
            handler_names = [handler_names]
        assert isinstance(handler_names, (list, tuple)) is True, handler_names

        exclude_handlers = exclude_handlers or []
        handlers = []
        for name in handler_names:
            if name in logging._handlers:
                handler = logging._handlers[name]
                if handler not in exclude_handlers:
                    handlers.append(handler)
        return handlers