From 9db4f6b4430d558eee2e4b78e6026181e8213430 Mon Sep 17 00:00:00 2001
From: Alexander Beedie <alexander-beedie@users.noreply.github.com>
Date: Fri, 9 Jun 2023 09:13:10 +0400
Subject: [PATCH] feat(python): allow use of `Config` object as a function
 decorator (#9307)

---
 py-polars/docs/source/reference/config.rst | 18 ++++++++++--
 py-polars/polars/config.py                 | 14 +++++++--
 py-polars/tests/unit/test_cfg.py           | 33 +++++++++++++---------
 3 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/py-polars/docs/source/reference/config.rst b/py-polars/docs/source/reference/config.rst
index 4e78916618ca..123c5427dc60 100644
--- a/py-polars/docs/source/reference/config.rst
+++ b/py-polars/docs/source/reference/config.rst
@@ -55,10 +55,24 @@ explicitly calling one or more of the available "set\_" methods on it...
 
     # on scope exit any modified settings are restored to their previous state
 
-...or by setting the options in the ``Config`` init directly (optionally
-omitting the "set\_" prefix for brevity):
+...or, often cleaner, by setting the options in the ``Config`` init directly
+(optionally omitting the "set\_" prefix for brevity):
 
 .. code-block:: python
 
     with pl.Config(verbose=True):
         do_various_things()
+
+Use as a function decorator
+---------------------------
+
+In the same vein, you can also use ``Config`` as a function decorator to
+temporarily set options for the duration of the function call:
+
+.. code-block:: python
+
+    @pl.Config(set_ascii_tables=True)
+    def write_ascii_frame_to_stdout(df: pl.DataFrame) -> None:
+        sys.stdout.write(str(df))
+
+"""
diff --git a/py-polars/polars/config.py b/py-polars/polars/config.py
index 11094ed97b6c..4200682ddd15 100644
--- a/py-polars/polars/config.py
+++ b/py-polars/polars/config.py
@@ -57,14 +57,14 @@ def _get_float_fmt() -> str:
 _POLARS_CFG_DIRECT_VARS = {"set_fmt_float": _get_float_fmt}
 
 
-class Config:
+class Config(contextlib.ContextDecorator):
     """
     Configure polars; offers options for table formatting and more.
 
     Notes
     -----
-    Can also be used as a context manager in order to temporarily scope
-    the lifetime of specific options. For example:
+    Can also be used as a context manager OR a function decorator in order to
+    temporarily scope the lifetime of specific options. For example:
 
     >>> with pl.Config() as cfg:
     ...     # set verbose for more detailed output within the scope
@@ -80,6 +80,14 @@ class Config:
 
     (The compact format is available for all `Config` methods that take a single value).
 
+    Alternatively, you can use as a decorator in order to scope the duration of the
+    selected options to a specific function:
+
+    >>> @pl.Config(verbose=True)
+    ... def test():
+    ...     pass
+    ...
+
     """
 
     _original_state: str = ""
diff --git a/py-polars/tests/unit/test_cfg.py b/py-polars/tests/unit/test_cfg.py
index 6d9a9c354249..53aa472a8cd6 100644
--- a/py-polars/tests/unit/test_cfg.py
+++ b/py-polars/tests/unit/test_cfg.py
@@ -20,24 +20,25 @@ def _environ() -> Iterator[None]:
 def test_ascii_tables() -> None:
     df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
 
+    ascii_table_repr = (
+        "shape: (3, 3)\n"
+        "+-----+-----+-----+\n"
+        "| a   | b   | c   |\n"
+        "| --- | --- | --- |\n"
+        "| i64 | i64 | i64 |\n"
+        "+=================+\n"
+        "| 1   | 4   | 7   |\n"
+        "| 2   | 5   | 8   |\n"
+        "| 3   | 6   | 9   |\n"
+        "+-----+-----+-----+"
+    )
     # note: expect to render ascii only within the given scope
     with pl.Config(set_ascii_tables=True):
-        assert (
-            str(df) == "shape: (3, 3)\n"
-            "+-----+-----+-----+\n"
-            "| a   | b   | c   |\n"
-            "| --- | --- | --- |\n"
-            "| i64 | i64 | i64 |\n"
-            "+=================+\n"
-            "| 1   | 4   | 7   |\n"
-            "| 2   | 5   | 8   |\n"
-            "| 3   | 6   | 9   |\n"
-            "+-----+-----+-----+"
-        )
+        assert repr(df) == ascii_table_repr
 
     # confirm back to utf8 default after scope-exit
     assert (
-        str(df) == "shape: (3, 3)\n"
+        repr(df) == "shape: (3, 3)\n"
         "┌─────┬─────┬─────┐\n"
         "│ a   ┆ b   ┆ c   │\n"
         "│ --- ┆ --- ┆ --- │\n"
@@ -49,6 +50,12 @@ def test_ascii_tables() -> None:
         "└─────┴─────┴─────┘"
     )
 
+    @pl.Config(set_ascii_tables=True)
+    def ascii_table() -> str:
+        return repr(df)
+
+    assert ascii_table() == ascii_table_repr
+
 
 def test_hide_header_elements() -> None:
     df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})