Coverage for muutils/collect_warnings.py: 0%

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