Temporary attribute context manager for factories

Summary: Utility context manager that allows us to temporarily modify the class attributes of a factory to make a module for a subclass. Now used in the pairwise factory.

Differential Revision: D74192292

Author: Jason Chow

aepsych/factory/pairwise.py

@@ -17,6 +17,7 @@
     default_mean_covar_factory,
     DefaultMeanCovarFactory,
 )
+from aepsych.factory.utils import temporary_attributes
 from aepsych.kernels.pairwisekernel import PairwiseKernel
 
 
@@ -84,30 +85,21 @@ def _make_covar_module(self) -> gpytorch.kernels.Kernel:
             )
 
         # Temporarily modify attributes to make base pair covariance module
-        original_dim = self.dim
-        self.dim = len(active_dims) // 2
-        self.stimuli_per_trial = 1
-
-        base_cov = super()._make_covar_module()  # TODO: This is really awkward
-
-        self.dim = original_dim
-        self.stimuli_per_trial = 1
+        with temporary_attributes(self, dim=len(active_dims) // 2, stimuli_per_trial=1):
+            base_cov = super()._make_covar_module()
 
         if len(self.shared_dims) == 0:
             return PairwiseKernel(base_cov)
 
         else:  # Some paired dims
-            # Need to make an extra shared dim covariance module
-            self.dim = len(self.shared_dims)
-            orig_active_dims = self.active_dims
-            self.active_dims = self.shared_dims
-            self.stimuli_per_trial = 1
-
-            shared_cov = super()._make_covar_module()
-
-            self.dim = original_dim
-            self.stimuli_per_trial = 1
-            self.active_dims = orig_active_dims
+            # Again temporary attributes
+            with temporary_attributes(
+                self,
+                dim=len(self.shared_dims),
+                active_dims=self.shared_dims,
+                stimuli_per_trial=1,
+            ):
+                shared_cov = super()._make_covar_module()
 
             return PairwiseKernel(base_cov, active_dims=active_dims) * shared_cov
 

aepsych/factory/utils.py

@@ -4,9 +4,24 @@
 
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
-
+from contextlib import contextmanager
 
 __default_invgamma_concentration = 4.6
 __default_invgamma_rate = 1.0
 DEFAULT_INVGAMMA_CONC = 4.6
 DEFAULT_INVGAMMA_RATE = 1.0
+
+
+@contextmanager
+def temporary_attributes(obj, **kwargs):
+    """Temporarily sets attributes on an object, and restores them when the context exits."""
+
+    try:
+        old_attrs = {}
+        for attr, val in kwargs.items():
+            old_attrs[attr] = getattr(obj, attr)
+            setattr(obj, attr, val)
+        yield obj
+    finally:
+        for attr, val in old_attrs.items():
+            setattr(obj, attr, val)