Non-blocking exception catching context manager with step skipping and lightweight logging and reporting capabilities.

failstate.py

"""Non-blocking exception catching context manager with step skipping and
lightweight logging and reporting capabilities.

2021 - Andras Fulop @fulibacsi"""

import json
import sys
import traceback
import warnings


class SkipWithBlock(Exception):
    """Custom Exception for skipping code block execution."""


class StateFailed(Exception):
    """General exception for FAILSTATE fails."""


class FAILSTATE:
    """Non-blocking exception catcher context manager with step skipping
    and lightweight logging and reporting capabilities.

    External logger is also supported, in case no external logger is set, print
    statements and warning.warn calls will be used. External logger can be set
    at any point.

    Example usage:
    ```
    state = FAILSTATE(['case1', 'case2', 'case3'])

    with state('case1'):
        random_code_block
    >>> executed normally

    with state('case2'):
        random_failing_code_block
    >>> exception and traceback printed to stdout

    with state('case3'):
        random_code_block
    >>> warning that the code block is skipped

    state.generate_report()
    >>> {'case1': 'true',
    >>>  'case2': 'false',
    >>>  'case3': 'null',
    >>>  'success': 'false'}
    ```
    """

    def __init__(self, steps, logger=None):
        """Initialize the context manager.

        Params:
        -------
        steps: iterable
            Name of the steps / code blocks to keep track
        logger: logging object (default: None)
            External logger to use
        """
        self.FAILSTATE = False
        self.steps = {step: None for step in steps}
        self.set_logger(logger)

    def __call__(self, step):
        """Prepares context manager to work on a predifined step / code block.
        Raises error if step name is not from the predefined step list.

        Params:
        -------
        step: str
            One of the predefined step name

        Returns:
        --------
        self
        """
        if self.FAILSTATE:
            self.log_warning("Skipping execution as previous entries "
                             "have already failed.")

        if step not in self.steps:
            raise ValueError(f"Invalid step name {step}! "
                             f"Available steps are: "
                             f"{','.join(list(self.steps))}")
        self.actual_step = step
        return self

    def __enter__(self):
        """Preparing environment for execution.
        Code block skipping is set up here with a special exception class
        raised by a tracing function and captured specifically in __exit__.

        More details on the block skipping:
        https://code.google.com/archive/p/ouspg/wikis/AnonymousBlocksInPython.wiki
        """
        if self.FAILSTATE:
            sys.settrace(lambda *args, **keys: None)
            frame = sys._getframe(1)
            frame.f_trace = self.trace
        return self

    def trace(self, frame, event, arg):
        """Special tracing function to raise the SkipWithBlock exception."""
        raise SkipWithBlock()

    def __exit__(self, exception_type, exception_value, tb):
        """Wrapping up code block execution, updating success/fail state."""
        # Handling successful state
        if exception_type is None:
            self.steps[self.actual_step] = True

        # Handling block skipping
        elif issubclass(exception_type, SkipWithBlock):
            pass

        # Managing failed state
        elif exception_type is not None:
            self.FAILSTATE = True
            self.steps[self.actual_step] = False

            formatted_tb = '\n'.join(traceback.format_tb(tb))
            self.log_error(f"Execution failed with the following error: "
                           f"{exception_value}\n"
                           f"Traceback:\n{formatted_tb}")

        # suppress any raised exception
        return True

    def reset(self):
        """Reset to initial state."""
        self.FAILSTATE = False
        self.steps = {step: None for step in self.steps}

    def set_logger(self, logger=None):
        """Setting up logger to use.
        In case the logger is not set, built-in print and warnings.warn calls
        will be used.

        Params:
        -------
        logger: logging object or None (default: None)
            Logger object to use
        """
        self.logger = logger
        self.log_info = print if logger is None else logger.info
        self.log_warning = warnings.warn if logger is None else logger.warning
        self.log_error = warnings.warn if logger is None else logger.error

    def generate_report(self):
        """Generate report about the current state.
        A json string is generated from the results with a final success value.
        The values in the json string are:
            - success: 'true'
            - failed: 'false'
            - skipped / not executed: 'null'

        Returns:
        --------
        report: str
            JSON string containing the run results
        """
        self.steps['success'] = not self.FAILSTATE
        return json.dumps(self.steps)

    def log_results(self, logger=None):
        """Log the run results one-by-one with the provided logger.

        If no logger was set previously and a logger is specified during the
        call, set it as the logger before the actual result logging.

        The logger argument to replace the default logger is presented here 
        in case the logger object was not available during the initialization 
        of the object.

        Params:
        -------
        logger: logging object or None (default: None)
            The logging object to set before logging,
            if not set, use the actual logger
        """
        if self.logger is None and logger is not None:
            self.set_logger(logger)

        for step, success in self.steps.items():
            if success is False:
                self.log_error(f"Check {step} failed.")
            elif success is None:
                self.log_warning(f"Check {step} skipped.")
            else:
                self.log_info(f"Check {step} succeeded.")

    def raise_for_failure(self):
        """Raise an exception with failed step count."""
        if self.FAILSTATE:
            n_fails = sum(1 for step in self.steps.values() if step is False)
            n_steps = len(self.steps)
            percentage = n_fails / n_steps

            raise StateFailed(f"{percentage:4.0%} ({n_fails} / {n_steps})"
                              f" steps failed.")