FEAT-#7141: Add an ability to use config variables with a context manager

docs/flow/modin/config.rst

@@ -56,3 +56,30 @@ API.
     # Changing value of `NPartitions`
     modin.config.NPartitions.put(16)
     print(modin.config.NPartitions.get()) # prints '16'
+
+One can also use config variables with a context manager in order to use
+some config only for a certain part of the code:
+
+.. code-block:: python
+
+    import modin.config as cfg
+
+    # Default value for this config is 'False'
+    print(cfg.RangePartitioning.get()) # False
+
+    # Set the config to 'True' inside of the context-manager
+    with cfg.context(RangePartitioning=True):
+        print(cfg.RangePartitioning.get()) # True
+        df.merge(...) # will use range-partitioning impl
+
+    # Once the context is over, the config gets back to its previous value
+    print(cfg.RangePartitioning.get()) # False
+
+    # You can also set multiple config at once when you pass a dictionary to 'cfg.context'
+    print(cfg.AsyncReadMode.get()) # False
+
+    with cfg.context(RangePartitioning=True, AsyncReadMode=True):
+        print(cfg.RangePartitioning.get()) # True
+        print(cfg.AsyncReadMode.get()) # True
+    print(cfg.RangePartitioning.get()) # False
+    print(cfg.AsyncReadMode.get()) # False

modin/config/__init__.py

@@ -56,12 +56,13 @@
     TestReadFromSqlServer,
     TrackFileLeaks,
 )
-from modin.config.pubsub import Parameter, ValueSource
+from modin.config.pubsub import Parameter, ValueSource, context
 
 __all__ = [
     "EnvironmentVariable",
     "Parameter",
     "ValueSource",
+    "context",
     # General settings
     "IsDebug",
     "Engine",

modin/config/pubsub.py

@@ -13,6 +13,7 @@
 
 """Module houses ``Parameter`` class - base class for all configs."""
 
+import contextlib
 import warnings
 from collections import defaultdict
 from enum import IntEnum
@@ -21,6 +22,7 @@
     Any,
     Callable,
     DefaultDict,
+    Iterator,
     NamedTuple,
     Optional,
     Tuple,
@@ -449,4 +451,49 @@ def add_option(cls, choice: Any) -> Any:
         raise TypeError("Cannot add a choice to a parameter where choices is None")
 
 
-__all__ = ["Parameter"]
+@contextlib.contextmanager
+def context(**config: dict[str, Any]) -> Iterator[None]:
+    """
+    Set a value(s) for the specified config(s) from ``modin.config`` in the scope of the context.
+
+    Parameters
+    ----------
+    **config : dict[Parameter, Any]
+        Keyword describing a name of a config variable from ``modin.config`` as a key
+        and a new value as a value.
+
+    Examples
+    --------
+    >>> RangePartitioning.get()
+    False
+    >>> with context(RangePartitioning=True):
+    ...     print(RangePartitioning.get()) # True
+    True
+    False
+    >>> RangePartitioning.get()
+    False
+    >>> with context(RangePartitioning=True, AsyncReadMode=True):
+    ...     print(RangePartitioning.get()) # True
+    ...     print(AsyncReadMode.get()) # True
+    True
+    True
+    >>> RangePartitioning.get()
+    False
+    >>> AsyncReadMode.get()
+    False
+    """
+    import modin.config as cfg
+
+    old_values = {}
+    for name, val in config.items():
+        var = getattr(cfg, name)
+        old_values[var] = var.get()
+        var.put(val)
+    try:
+        yield
+    finally:
+        for var, val in old_values.items():
+            var.put(val)
+
+
+__all__ = ["Parameter", "context"]

modin/config/test/test_envvars.py

@@ -298,3 +298,61 @@ def get_values():
     finally:
         deprecated_var.put(old_depr_val)
         new_var.put(old_new_var)
+
+
+@pytest.mark.parametrize(
+    "modify_config",
+    [{cfg.RangePartitioning: False, cfg.LazyExecution: "Auto"}],
+    indirect=True,
+)
+def test_context_manager_update_config(modify_config):
+    # simple case, 1 parameter
+    assert cfg.RangePartitioning.get() is False
+    with cfg.context(RangePartitioning=True):
+        assert cfg.RangePartitioning.get() is True
+    assert cfg.RangePartitioning.get() is False
+
+    # nested case, 1 parameter
+    assert cfg.RangePartitioning.get() is False
+    with cfg.context(RangePartitioning=True):
+        assert cfg.RangePartitioning.get() is True
+        with cfg.context(RangePartitioning=False):
+            assert cfg.RangePartitioning.get() is False
+            with cfg.context(RangePartitioning=False):
+                assert cfg.RangePartitioning.get() is False
+            assert cfg.RangePartitioning.get() is False
+        assert cfg.RangePartitioning.get() is True
+    assert cfg.RangePartitioning.get() is False
+
+    # simple case, 2 parameters
+    assert cfg.RangePartitioning.get() is False
+    assert cfg.LazyExecution.get() == "Auto"
+    with cfg.context(RangePartitioning=True, LazyExecution="Off"):
+        assert cfg.RangePartitioning.get() is True
+        assert cfg.LazyExecution.get() == "Off"
+    assert cfg.RangePartitioning.get() is False
+    assert cfg.LazyExecution.get() == "Auto"
+
+    # nested case, 2 parameters
+    assert cfg.RangePartitioning.get() is False
+    assert cfg.LazyExecution.get() == "Auto"
+    with cfg.context(RangePartitioning=True, LazyExecution="Off"):
+        assert cfg.RangePartitioning.get() is True
+        assert cfg.LazyExecution.get() == "Off"
+        with cfg.context(RangePartitioning=False):
+            assert cfg.RangePartitioning.get() is False
+            assert cfg.LazyExecution.get() == "Off"
+            with cfg.context(LazyExecution="On"):
+                assert cfg.RangePartitioning.get() is False
+                assert cfg.LazyExecution.get() == "On"
+                with cfg.context(RangePartitioning=True, LazyExecution="Off"):
+                    assert cfg.RangePartitioning.get() is True
+                    assert cfg.LazyExecution.get() == "Off"
+                assert cfg.RangePartitioning.get() is False
+                assert cfg.LazyExecution.get() == "On"
+            assert cfg.RangePartitioning.get() is False
+            assert cfg.LazyExecution.get() == "Off"
+        assert cfg.RangePartitioning.get() is True
+        assert cfg.LazyExecution.get() == "Off"
+    assert cfg.RangePartitioning.get() is False
+    assert cfg.LazyExecution.get() == "Auto"