Merge pull request #534 from gchq/fix/remaining-decorator-warnings

fix: decorator warnings raised by test_basedecorator.py

Author: pc532627

tests/units/decoratorutils/test_basedecorator.py

@@ -18,11 +18,13 @@
 from typing import Any, TypeVar, Union
 
 import pytest
+from typing_extensions import ContextManager
 
 from tests.cases import assert_not_warns
 from vanguard.base import GPController
 from vanguard.decoratorutils import Decorator, TopMostDecorator, errors, wraps_class
 from vanguard.decoratorutils.errors import TopmostDecoratorError
+from vanguard.utils import multi_context
 from vanguard.vanilla import GaussianGPController
 
 ControllerT = TypeVar("ControllerT", bound=GPController)
@@ -53,7 +55,7 @@ def __init__(self, **kwargs: Any) -> None:
         super().__init__(framework_class=SimpleNumber, required_decorators={}, **kwargs)
 
     def _decorate_class(self, cls: type[SimpleNumber]) -> type[SimpleNumber]:
-        @wraps_class(cls)
+        @wraps_class(cls, decorator_source=self)
         class InnerClass(cls):
             """A wrapper for normalising y inputs and variance."""
 
@@ -319,29 +321,40 @@ def add_5(self) -> Union[int, float]:
 
             # If we're not ignoring the errors, then specify what the expected error/warning is.
             if who_is_wrong == "subclass":
-                blame_class = "NewNumber"
+                blame_classes = ["NewNumber"]
+            elif who_is_wrong == "superclass":
+                blame_classes = ["MiddleNumber"]
             elif who_is_wrong == "superclass" or who_is_wrong == "both":
-                blame_class = "MiddleNumber"
+                blame_classes = ["MiddleNumber", "NewNumber"]
             else:
                 raise ValueError(who_is_wrong)
 
+            contexts: list[ContextManager]
             if mode == "new_method":
-                expected_message = (
+                expected_messages = [
                     f"{SquareResult.__name__!r}: The class '{blame_class}' has added the following unexpected methods"
+                    for blame_class in blame_classes
+                ]
+                contexts = (
+                    [pytest.raises(errors.UnexpectedMethodError, match=expected_messages[0])]
+                    if raise_instead
+                    else [pytest.warns(errors.UnexpectedMethodWarning, match=msg) for msg in expected_messages]
                 )
-                expected_types = errors.UnexpectedMethodError, errors.UnexpectedMethodWarning
             elif mode == "override_method":
-                expected_message = (
+                expected_messages = [
                     f"{SquareResult.__name__!r}: The class '{blame_class}' has overwritten the following methods"
+                    for blame_class in blame_classes
+                ]
+                contexts = (
+                    [pytest.raises(errors.OverwrittenMethodError, match=expected_messages[0])]
+                    if raise_instead
+                    else [pytest.warns(errors.OverwrittenMethodWarning, match=msg) for msg in expected_messages]
                 )
-                expected_types = errors.OverwrittenMethodError, errors.OverwrittenMethodWarning
             else:
                 raise ValueError(mode)
 
-            if raise_instead:
-                context = pytest.raises(expected_types[0], match=expected_message)
-            else:
-                context = pytest.warns(expected_types[1], match=expected_message)
+            context = multi_context(contexts)
+
         elif ignore == "specific":
             # Ignore by setting `ignore_methods`.
             if mode == "new_method":
@@ -358,9 +371,11 @@ def add_5(self) -> Union[int, float]:
         else:
             raise ValueError(ignore)
 
-        # Actually perform the decoration, and check for errors/warnings as appropriate
-        with context:
-            SquareResult(**kwargs)(NewNumber)
+        # Actually perform the decoration, and check for errors/warnings as appropriate. No warnings other than those
+        # we expect should be raised.
+        with assert_not_warns(errors.OverwrittenMethodWarning, errors.UnexpectedMethodWarning):
+            with context:
+                SquareResult(**kwargs)(NewNumber)
 
 
 class TestInvalidDecoration:
@@ -394,7 +409,7 @@ class NewNumber(SimpleNumber):  # pylint: disable=unused-variable
     def test_passed_requirements(self) -> None:
         """Test that when all requirements for a decorator are satisfied, no error is thrown."""
 
-        @RequiresSquareResult(ignore_methods=("__init__", "add_5"))
+        @RequiresSquareResult(ignore_all=True)
         @SquareResult()
         class NewNumber(SimpleNumber):  # pylint: disable=unused-variable
             """Declaring this class should not throw any error, as all requirements are satisfied."""

tests/units/test_decorator_combinations.py

@@ -379,10 +379,18 @@ def _decorate_class(self, cls: T) -> T:
 # (upper, lower) -> (warning type, message regex)
 # Warnings we expect to be raised on decorator application.
 EXPECTED_COMBINATION_APPLY_WARNINGS: dict[tuple[type[Decorator], type[Decorator]], tuple[type[Warning], str]] = {
-    (NormaliseY, DirichletMulticlassClassification): (
-        BadCombinationWarning,
-        "NormaliseY should not be used above classification decorators - this may lead to unexpected behaviour.",
-    ),
+    **{
+        (NormaliseY, lower): (
+            BadCombinationWarning,
+            "NormaliseY should not be used above classification decorators - this may lead to unexpected behaviour.",
+        )
+        for lower in [
+            BinaryClassification,
+            CategoricalClassification,
+            DirichletMulticlassClassification,
+            DirichletKernelMulticlassClassification,
+        ]
+    },
     **{
         (VariationalInference, lower): (
             BadCombinationWarning,

tests/units/test_utils.py

@@ -17,11 +17,13 @@
 """
 
 import unittest
+from unittest.mock import MagicMock
 
 import numpy as np
 import numpy.typing
 import pytest
 import torch
+from typing_extensions import ContextManager
 
 from tests.cases import get_default_rng
 from vanguard.utils import (
@@ -30,6 +32,7 @@
     generator_append_constant,
     infinite_tensor_generator,
     instantiate_with_subset_of_kwargs,
+    multi_context,
     optional_random_generator,
 )
 
@@ -286,3 +289,37 @@ def test_generator_zero_dimensional(self):
 
         with pytest.raises(ValueError, match="0-dimensional tensors are incompatible"):
             next(generator)
+
+
+@pytest.mark.parametrize("num_contexts", [1, 5])
+def test_multi_context(num_contexts: int):
+    """
+    Test the `multi_context` context manager.
+
+    Given that the context manager syntax is ultimately just syntactic sugar for calling `ctx.__enter__` and
+    `ctx.__exit__` (with some extra semantics around exception handling), all we really need to do is check that
+    these methods are called on each context manager passed to `multi_context` at the appropriate times. We don't
+    test the specifics of exception propagation, as that is (a) too complex for a single unit test, and (b) more the
+    responsibility of the `contextlib.ExitStack` that `multi_context` uses internally.
+
+    Test with both a single context and multiple contexts being passed to `multi_context`.
+    """
+    dummy_contexts = [MagicMock(spec=ContextManager)() for _ in range(num_contexts)]
+    context = multi_context(dummy_contexts)
+
+    for ctx in dummy_contexts:
+        # check that we haven't entered or left any of the contexts yet
+        ctx.__enter__.assert_not_called()
+        ctx.__exit__.assert_not_called()
+
+    # enter the multi context
+    with context:
+        for ctx in dummy_contexts:
+            # check that we have entered but not left each of the contexts yet
+            ctx.__enter__.assert_called_once()
+            ctx.__exit__.assert_not_called()
+
+    for ctx in dummy_contexts:
+        # check that we have left each of the contexts
+        ctx.__enter__.assert_called_once()
+        ctx.__exit__.assert_called_once()

vanguard/base/gpcontroller.py

@@ -141,7 +141,8 @@ def fit(
                 warnings.warn(
                     f"You are trying to set gradient_every (in this case to {gradient_every}) in batch mode."
                     "This does not make mathematical sense and your value of gradient every will be ignored "
-                    " and replaced by 1."
+                    " and replaced by 1.",
+                    stacklevel=2,
                 )
             gradient_every = 1
 

vanguard/decoratorutils/basedecorator.py

@@ -185,6 +185,9 @@ def _verify_class_has_no_newly_added_methods(self, cls: type[T], super_methods:
         cls_methods = {
             key
             for key, value in getmembers(cls, isfunction)
+            # only functions that are actually from this class (as opposed to a superclass)
+            if key in cls.__dict__
+            # ignore functions defined in safe_updates
             if key not in self.safe_updates.get(self._get_method_implementation(cls, key), set())
             # beartype does weird things with __sizeof__; however, it's of no concern to us, and we never make use of
             # this dunder attribute. See https://github.com/beartype/beartype/blob/v0.19.0/beartype/_decor/_decortype.py
@@ -216,7 +219,7 @@ def _verify_class_has_no_newly_added_methods(self, cls: type[T], super_methods:
             else:
                 warnings.warn(message, errors.UnexpectedMethodWarning, stacklevel=4)
 
-        overwritten_methods = {method for method in cls_methods if method in cls.__dict__} - ignore_methods
+        overwritten_methods = cls_methods - ignore_methods - extra_methods
         if overwritten_methods:
             if __debug__:
                 overwritten_method_messages = "\n".join(

vanguard/distribute/decorator.py

@@ -275,7 +275,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                     **self._expert_init_kwargs,
                 )
 
-            def fit(self, n_sgd_iters: int = 10, gradient_every: int = 10) -> torch.Tensor:
+            def fit(self, n_sgd_iters: int = 10, gradient_every: Optional[int] = None) -> torch.Tensor:
                 """
                 Create the expert controllers.
 

vanguard/utils.py

@@ -16,15 +16,17 @@
 Contain some small utilities of use in some cases.
 """
 
+import contextlib
 import functools
 import os
 import warnings
-from collections.abc import Generator
+from collections.abc import Generator, Iterable
 from typing import Any, Callable, Optional, TypeVar
 
 import numpy as np
 import numpy.typing
 import torch
+from typing_extensions import ContextManager
 
 from vanguard.warnings import _RE_INCORRECT_LIKELIHOOD_PARAMETER
 
@@ -253,3 +255,12 @@ def compose(functions: list[Callable[[T], T]]) -> Callable[[T], T]:
     :return: A single function of type (T -> T) that applies each of the passed functions in series.
     """
     return lambda x: functools.reduce(lambda acc, f: f(acc), reversed(functions), x)
+
+
+@contextlib.contextmanager
+def multi_context(contexts: Iterable[ContextManager]):
+    """Combine multiple context managers into one."""
+    with contextlib.ExitStack() as stack:
+        for context in contexts:
+            stack.enter_context(context)
+        yield