add context manager functionality to config.defaults

murrayrm · murrayrm · commit b08321268acc · 2025-01-12T17:28:33.000-08:00
diff --git a/control/config.py b/control/config.py
@@ -73,6 +73,28 @@ def _check_deprecation(self, key):
         else:
             return key
 
+    #
+    # Context manager functionality
+    #
+
+    def __call__(self, mapping):
+        self.saved_mapping = dict()
+        self.temp_mapping = mapping.copy()
+        return self
+
+    def __enter__(self):
+        for key, val in self.temp_mapping.items():
+            if not key in self:
+                raise ValueError(f"unknown parameter '{key}'")
+            self.saved_mapping[key] = self[key]
+            self[key] = val
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for key, val in self.saved_mapping.items():
+            self[key] = val
+        del self.saved_mapping, self.temp_mapping
+        return None
 
 defaults = DefaultDict(_control_defaults)
 
diff --git a/control/tests/config_test.py b/control/tests/config_test.py
@@ -332,3 +332,20 @@ def test_legacy_repr_format(self):
         new = eval(repr(sys))
         for attr in ['A', 'B', 'C', 'D']:
             assert getattr(sys, attr) == getattr(sys, attr)
+
+
+def test_config_context_manager():
+    # Make sure we can temporarily set the value of a parameter
+    default_val = ct.config.defaults['statesp.latex_repr_type']
+    with ct.config.defaults({'statesp.latex_repr_type': 'new value'}):
+        assert ct.config.defaults['statesp.latex_repr_type'] != default_val
+        assert ct.config.defaults['statesp.latex_repr_type'] == 'new value'
+    assert ct.config.defaults['statesp.latex_repr_type'] == default_val
+
+    # OK to call the context manager and not do anything with it
+    ct.config.defaults({'statesp.latex_repr_type': 'new value'})
+    assert ct.config.defaults['statesp.latex_repr_type'] == default_val
+
+    with pytest.raises(ValueError, match="unknown parameter 'unknown'"):
+        with ct.config.defaults({'unknown': 'new value'}):
+            pass
