updated options_context and simplified dictionary handling in set_option

Author: arthurlw

pandas/_config/config.py

@@ -264,18 +264,7 @@ def set_option(*args) -> None:
     """
     # Handle dictionary input
     if len(args) == 1 and isinstance(args[0], dict):
-        options_dict = args[0]
-        for k, v in options_dict.items():
-            key = _get_single_key(k)
-            opt = _get_registered_option(key)
-            if opt and opt.validator:
-                opt.validator(v)
-            # walk the nested dict
-            root, k_root = _get_root(key)
-            root[k_root] = v
-            if opt.cb:
-                opt.cb(key)
-        return
+        args = tuple(kv for item in args[0].items() for kv in item)
 
     # Handle single option-value pair
     if len(args) == 2:
@@ -485,9 +474,10 @@ def option_context(*args) -> Generator[None]:
 
     Parameters
     ----------
-    *args : str | object
+    *args : str | object | dict
         An even amount of arguments provided in pairs which will be
-        interpreted as (pattern, value) pairs.
+        interpreted as (pattern, value) pairs. Alternatively, a single
+        dictionary of {pattern: value} may be provided.
 
     Returns
     -------
@@ -516,7 +506,12 @@ def option_context(*args) -> Generator[None]:
     >>> from pandas import option_context
     >>> with option_context("display.max_rows", 10, "display.max_columns", 5):
     ...     pass
+    >>> with option_context({"display.max_rows": 10, "display.max_columns": 5}):
+    ...     pass
     """
+    if len(args) == 1 and isinstance(args[0], dict):
+        args = tuple(kv for item in args[0].items() for kv in item)
+
     if len(args) % 2 != 0 or len(args) < 2:
         raise ValueError(
             "Provide an even amount of arguments as "

pandas/tests/config/test_config.py

@@ -395,6 +395,34 @@ def f():
 
         f()
 
+    def test_set_ContextManager_dict(self):
+        def eq(val):
+            assert cf.get_option("a") == val
+            assert cf.get_option("b.c") == val
+
+        cf.register_option("a", 0)
+        cf.register_option("b.c", 0)
+
+        eq(0)
+        with cf.option_context({"a": 15, "b.c": 15}):
+            eq(15)
+            with cf.option_context({"a": 25, "b.c": 25}):
+                eq(25)
+            eq(15)
+        eq(0)
+
+        cf.set_option("a", 17)
+        cf.set_option("b.c", 17)
+        eq(17)
+
+        # Test that option_context can be used as a decorator too
+        @cf.option_context({"a": 123, "b.c": 123})
+        def f():
+            eq(123)
+
+        f()
+
+
     def test_attribute_access(self):
         holder = []