1from __future__ import annotations
  2
  3import sys
  4import warnings
  5from collections import Counter
  6from contextlib import AbstractContextManager
  7from types import TracebackType
  8from typing import Any, Literal
  9
 10
 11class CollateWarnings(AbstractContextManager["CollateWarnings"]):
 12    """Capture every warning issued inside a `with` block and print a collated
 13    summary when the block exits.
 14
 15    Internally this wraps `warnings.catch_warnings(record=True)` so that all
 16    warnings raised in the block are recorded.  When the context exits, identical
 17    warnings are grouped and (optionally) printed with a user-defined format.
 18
 19    # Parameters:
 20     - `print_on_exit : bool`
 21       Whether to print the summary when the context exits
 22       (defaults to `True`)
 23     - `fmt : str`
 24       Format string used for printing each line of the summary.
 25       Available fields are:
 26
 27       * `{count}`     : number of occurrences
 28       * `{filename}`  : file where the warning originated
 29       * `{lineno}`    : line number
 30       * `{category}`  : warning class name
 31       * `{message}`   : warning message text
 32
 33       (defaults to `"({count}x) {filename}:{lineno} {category}: {message}"`)
 34
 35    # Returns:
 36     - `CollateWarnings`
 37       The context-manager instance.  After exit, the attribute
 38       `counts` holds a mapping
 39
 40       ```python
 41       {(filename, lineno, category, message): count}
 42       ```
 43
 44    # Usage:
 45    ```python
 46    >>> import warnings
 47    >>> with CollateWarnings() as cw:
 48    ...     warnings.warn("deprecated", DeprecationWarning)
 49    ...     warnings.warn("deprecated", DeprecationWarning)
 50    ...     warnings.warn("other", UserWarning)
 51    (2x) /tmp/example.py:42 DeprecationWarning: deprecated
 52    (1x) /tmp/example.py:43 UserWarning: other
 53    >>> cw.counts
 54    {('/tmp/example.py', 42, 'DeprecationWarning', 'deprecated'): 2,
 55     ('/tmp/example.py', 43, 'UserWarning', 'other'): 1}
 56    ```
 57    """
 58
 59    _active: bool
 60    _catcher: Any
 61    _records: list[warnings.WarningMessage]
 62    counts: Counter[
 63        tuple[
 64            str,  # filename
 65            int,  # lineno
 66            str,  # category name
 67            str,  # message
 68        ]
 69    ]
 70    print_on_exit: bool
 71    fmt: str
 72
 73    def __init__(
 74        self,
 75        print_on_exit: bool = True,
 76        fmt: str = "({count}x) {filename}:{lineno} {category}: {message}",
 77    ) -> None:
 78        self.print_on_exit = print_on_exit
 79        self.fmt = fmt
 80        self._active = False
 81        self._records = []
 82        self.counts = Counter()
 83
 84    def __enter__(self) -> CollateWarnings:
 85        if self._active:
 86            raise RuntimeError("CollateWarnings cannot be re-entered")
 87
 88        self._active = True
 89        self._catcher = warnings.catch_warnings(record=True)
 90        self._records = self._catcher.__enter__()
 91        warnings.simplefilter("always")  # capture every warning
 92        return self
 93
 94    def __exit__(
 95        self,
 96        exc_type: type[BaseException] | None,
 97        exc_val: BaseException | None,
 98        exc_tb: TracebackType | None,
 99    ) -> Literal[False]:
100        if not self._active:
101            raise RuntimeError("CollateWarnings exited twice")
102
103        self._active = False
104        # stop capturing
105        self._catcher.__exit__(exc_type, exc_val, exc_tb)
106
107        # collate
108        self.counts = Counter(
109            (
110                rec.filename,
111                rec.lineno,
112                rec.category.__name__,
113                str(rec.message),
114            )
115            for rec in self._records
116        )
117
118        if self.print_on_exit:
119            for (filename, lineno, category, message), count in self.counts.items():
120                print(
121                    self.fmt.format(
122                        count=count,
123                        filename=filename,
124                        lineno=lineno,
125                        category=category,
126                        message=message,
127                    ),
128                    file=sys.stderr,
129                )
130
131        # propagate any exception from the with-block
132        return False