Source code for schrodinger.test.pytest.warnings

"""
Turn existing python warnings into failures, for use in pytest. Uses the
internals of the cpython warning module to modify existing warning filters to
change warnings.filterwarnings action into errors.
"""

import copy
import dataclasses
import re
import warnings
from typing import Optional


@dataclasses.dataclass
class _WarningFilter:
    action: str
    message: Optional[re.Pattern]
    category: Warning
    module: Optional[re.Pattern]
    lineno: int

    def as_filterwarnings_args(self):
        new_filter = copy.copy(self)
        new_filter.message = _munge_regex_to_str(new_filter.message)
        new_filter.module = _munge_regex_to_str(new_filter.module)
        return dataclasses.astuple(new_filter)


def _get_new_warning_filter(warning_filter) -> _WarningFilter:
    """
    Modify an existing warning filter, to change modes to error from printing.
    """
    new_filter = _WarningFilter(*warning_filter)
    if new_filter.action in {"always", "default", "module"}:
        new_filter.action = "error"
    return new_filter.as_filterwarnings_args()


def mark_warnings_as_errors():
    """
    Modify existing warning filters using python internals to mark warnings
    as errors, causing them to raise as pytest errors. Does not apply to
    subprocesses.
    """
    existing_filters = warnings.filters.copy()
    warnings.resetwarnings()
    for existing_warning in reversed(existing_filters):
        warnings.filterwarnings(*_get_new_warning_filter(existing_warning))


def _munge_regex_to_str(regex: Optional[re.Pattern]) -> str:
    """
    Helper function to convert compiled regex objects back to str, for use
    in reconstructing python warning objects.
    """
    if regex:
        try:
            return regex.pattern
        except AttributeError:
            if isinstance(regex, str):
                return regex
            raise
    return ""