Simplify _ChangeFlagDecorator

Author: Armavica

pytensor/configparser.py

@@ -32,11 +32,7 @@ class ConfigAccessViolation(AttributeError):
 
 
 class _ChangeFlagsDecorator:
-    def __init__(self, *args, _root=None, **kwargs):
-        # the old API supported passing a dict as the first argument:
-        if args:
-            assert len(args) == 1 and isinstance(args[0], dict)
-            kwargs = dict(**args[0], **kwargs)
+    def __init__(self, _root=None, **kwargs):
         self.confs = {k: _root._config_var_dict[k] for k in kwargs}
         self.new_vals = kwargs
         self._root = _root
@@ -310,14 +306,14 @@ def fetch_val_for_key(self, key, delete_key: bool = False):
         except (NoOptionError, NoSectionError):
             raise KeyError(key)
 
-    def change_flags(self, *args, **kwargs) -> _ChangeFlagsDecorator:
+    def change_flags(self, **kwargs) -> _ChangeFlagsDecorator:
         """
         Use this as a decorator or context manager to change the value of
         PyTensor config variables.
 
         Useful during tests.
         """
-        return _ChangeFlagsDecorator(*args, _root=self, **kwargs)
+        return _ChangeFlagsDecorator(_root=self, **kwargs)
 
     def warn_unused_flags(self):
         for key in self._flags_dict:

tests/link/c/test_type.py

@@ -287,6 +287,6 @@ def test_op_with_cenumtype(self):
         assert val_billion == val_million * 1000
         assert val_two_billions == val_billion * 2
 
-    @pytensor.config.change_flags(**{"cmodule__debug": True})
+    @pytensor.config.change_flags(cmodule__debug=True)
     def test_op_with_cenumtype_debug(self):
         self.test_op_with_cenumtype()

tests/tensor/test_blas.py

@@ -514,7 +514,7 @@ def compute_ref(
         C = self.get_value(C, transpose_C, slice_C)
         return alpha * np.dot(A, B) + beta * C
 
-    @config.change_flags({"blas__ldflags": ""})
+    @config.change_flags(blas__ldflags="")
     def run_gemm(
         self,
         dtype,

tests/test_config.py

@@ -168,7 +168,7 @@ def test_config_context():
 
     with root.change_flags(test__config_context="new_value"):
         assert root.test__config_context == "new_value"
-        with root.change_flags({"test__config_context": "new_value2"}):
+        with root.change_flags(test__config_context="new_value2"):
             assert root.test__config_context == "new_value2"
         assert root.test__config_context == "new_value"
     assert root.test__config_context == "test_default"