Merge pull request #662 from aai-institute/feature/better-warnings-filter

Allow suppress_warnings to raise exceptions / any other valid action

Author: mdbenito

CHANGELOG.md

@@ -58,8 +58,10 @@
 
 ### Fixed
 
-- Fixed `show_warnings=False` not being respected in subprocesses
+- Fixed `show_warnings=False` not being respected in subprocesses. Introduced
+  `suppress_warninigs` decorator for more flexibility
   [PR #647](https://github.com/aai-institute/pyDVL/pull/647)
+  [PR #662](https://github.com/aai-institute/pyDVL/pull/662)
 - Fixed several bugs in diverse stopping criteria, including: iteration counts,
   computing completion, resetting, nested composition
   [PR #641](https://github.com/aai-institute/pyDVL/pull/641)

src/pydvl/utils/functional.py

@@ -163,44 +163,63 @@ def suppress_warnings(
 ) -> Union[Callable[[Callable[P, R]], Callable[P, R]], Callable[P, R]]:
     """Decorator for class methods to conditionally suppress warnings.
 
-    The decorated method will execute with warnings suppressed for the specified
-    categories. If the instance has the attribute named by `flag`, and it evaluates to
-    `True`, then suppression will be deactivated.
-
-    ??? Example "Suppress all warnings"
-        ```python
-        class A:
-            @suppress_warnings
-            def method(self, ...):
-                ...
-        ```
-    ??? Example "Suppress only `UserWarning`"
-        ```python
-        class A:
-            @suppress_warnings(categories=(UserWarning,))
-            def method(self, ...):
-                ...
-        ```
-    ??? Example "Configuring behaviour at runtime"
-        ```python
-        class A:
-            def __init__(self, warn_enabled: bool):
-                self.warn_enabled = warn_enabled
-
-            @suppress_warnings(flag="warn_enabled")
-            def method(self, ...):
-                ...
-        ```
-
-    Args:
-        fun: Optional callable to decorate. If provided, the decorator is applied inline.
-        categories: Sequence of warning categories to suppress.
-        flag: Name of an instance attribute to check for enabling warnings. If the
-              attribute exists and evaluates to `True`, warnings will **not** be
-              suppressed.
-
-    Returns:
-        Either a decorator (if no function is provided) or the decorated callable.
+      The decorated method will execute with warnings suppressed for the specified
+      categories. If the instance has the attribute named by `flag`, and it's a boolean
+      evaluating to `False`, warnings will be ignored. If the attribute is a string, then
+      it is interpreted as an "action" to be performed on the categories specified.
+      Allowed values are as per [warnings.simplefilter][], which are:
+    `default`, `error`, `ignore`, `always`, `all`, `module`, `once`
+
+      ??? Example "Suppress all warnings"
+          ```python
+          class A:
+              @suppress_warnings
+              def method(self, ...):
+                  ...
+          ```
+      ??? Example "Suppress only `UserWarning`"
+          ```python
+          class A:
+              @suppress_warnings(categories=(UserWarning,))
+              def method(self, ...):
+                  ...
+          ```
+      ??? Example "Configuring behaviour at runtime"
+          ```python
+          class A:
+              def __init__(self, warn_enabled: bool):
+                  self.warn_enabled = warn_enabled
+
+              @suppress_warnings(flag="warn_enabled")
+              def method(self, ...):
+                  ...
+          ```
+
+      ??? Example "Raising on RuntimeWarning"
+          ```python
+          class A:
+              def __init__(self, warnings: str = "error"):
+                  self.warnings = warnings
+
+              @suppress_warnings(flag="warnings")
+              def method(self, ...):
+                  ...
+
+          A().method()  # Raises RuntimeWarning
+          ```
+
+
+      Args:
+          fun: Optional callable to decorate. If provided, the decorator is applied inline.
+          categories: Sequence of warning categories to suppress.
+          flag: Name of an instance attribute to check for enabling warnings. If the
+                attribute exists and evaluates to `False`, warnings will be ignored. If
+                it evaluates to a str, then this action will be performed on the categories
+                specified. Allowed values are as per [warnings.simplefilter][], which are:
+                `default`, `error`, `ignore`, `always`, `all`, `module`, `once`
+
+      Returns:
+          Either a decorator (if no function is provided) or the decorated callable.
     """
 
     def decorator(fn: Callable[P, R]) -> Callable[P, R]:
@@ -227,11 +246,18 @@ def wrapper(self, *args: Any, **kwargs: Any) -> R:
                     raise AttributeError(
                         f"Instance has no attribute '{flag}' for suppress_warnings"
                     )
-                if flag and getattr(self, flag, False):
+                if flag and getattr(self, flag, False) is True:
                     return fn(self, *args, **kwargs)
+                # flag is either False or a string
                 with warnings.catch_warnings():
+                    if (action := getattr(self, flag, "ignore")) is False:
+                        action = "ignore"
+                    elif not isinstance(action, str):
+                        raise TypeError(
+                            f"Expected a boolean or string for flag '{flag}', got {type(action).__name__}"
+                        )
                     for category in categories:
-                        warnings.simplefilter("ignore", category=category)
+                        warnings.simplefilter(action, category=category)  # type: ignore
                     return fn(self, *args, **kwargs)
 
             return cast(Callable[P, R], wrapper)

tests/utils/test_functional.py

@@ -1,4 +1,5 @@
-import gc
+from __future__ import annotations
+
 import inspect
 import logging
 import time
@@ -13,7 +14,7 @@
 
 
 class WarningsClass:
-    def __init__(self, show_warnings: bool = True):
+    def __init__(self, show_warnings: str | bool = True):
         self.show_warnings = show_warnings
 
     @suppress_warnings(categories=(UserWarning,), flag="show_warnings")
@@ -99,6 +100,18 @@ def test_different_categories(
     assert warning_message in str(w.message)
 
 
+def test_raises_on_flag_error():
+    obj = WarningsClass(show_warnings="error")
+    with pytest.raises(UserWarning):
+        obj.method_warn()
+
+
+def test_invalid_flag_type():
+    obj = WarningsClass(show_warnings=42)
+    with pytest.raises(TypeError):
+        obj.method_warn()
+
+
 def test_nonmethod_decorator_usage():
     @suppress_warnings(categories=(RuntimeWarning,))
     def fun(x: int) -> float: