Source code for mmf.utils.mmf_warnings

"""Enhanced warning facilities.

This may be used as a replacement for the standard warnings module.
"""
import sys
import warnings
import traceback

__all__ = warnings.__all__
from warnings import *
from warnings import simplefilter

# Set this to False if you do not want to print the warning stack
# trace to stderr
_warn_show_stack_trace = True


def warn(message,category=UserWarning,stacklevel=1,quiet=False,
         stdout_message=None):
    """Replacement for warnings.warn that also sends the stack trace
    to sys.stderr.  Supports additional argument stdout_message which
    is sent to stdout rather than stderr every time the warning is
    called.

    Example:
    >>> import mmf.utils.mmf_warnings as warnings
    >>> def f():
    ...    warnings.warn("Warning in f!",stdout_message="Short warning!")
    >>> def g():
    ...    f()
    >>> stderr = sys.stderr
    >>> sys.stderr = open('/dev/null','w') # Suppress stderr output
    >>> f()
    Short warning!
    >>> sys.stderr.close()
    >>> sys.stderr = sys.stdout # Send all output to stdout for checking.
    >>> g()                # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
    Short warning!
    ...: UserWarning: Warning in f!
      ...
      File "<doctest mmf.utils.mmf_warnings.warn[8]>", line 1, in <module>
        g()                # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
      File "<doctest mmf.utils.mmf_warnings.warn[2]>", line 2, in g
        f()
      File "<doctest mmf.utils.mmf_warnings.warn[1]>", line 2, in f
        warnings.warn("Warning in f!",stdout_message="Short warning!")
    >>> warnings._warn_show_stack_trace = False # You can turn off the trace
    >>> __warningregistry__ = {}
    >>> warnings.__warningregistry__ = {}
    >>> g()                # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
    Short warning!
    ...: UserWarning: Warning in f!
    """
    if _warn_show_stack_trace and not quiet:
        stack = traceback.extract_stack()
        format = traceback.format_list(stack)
        str = "".join(format[:-1])
        message = message + "\n" + str

    if stdout_message is not None:
        sys.stdout.write(stdout_message + "\n")

    warnings.warn(message,category,stacklevel=stacklevel+1)