gh-91291: Accept attributes as keyword arguments in decimal.localcontext (#32242)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

6 files changed, 161 insertions, 62 deletions

diff --git a/Doc/library/decimal.rst b/Doc/library/decimal.rst
index e759c5c..2ad84f2 100644
--- a/Doc/library/decimal.rst
+++ b/Doc/library/decimal.rst
@@ -925,12 +925,13 @@ Each thread has its own current context which is accessed or changed using the
 You can also use the :keyword:`with` statement and the :func:`localcontext`
 function to temporarily change the active context.
 
-.. function:: localcontext(ctx=None)
+.. function:: localcontext(ctx=None, \*\*kwargs)
 
    Return a context manager that will set the current context for the active thread
    to a copy of *ctx* on entry to the with-statement and restore the previous context
    when exiting the with-statement. If no context is specified, a copy of the
-   current context is used.
+   current context is used.  The *kwargs* argument is used to set the attributes
+   of the new context.
 
    For example, the following code sets the current decimal precision to 42 places,
    performs a calculation, and then automatically restores the previous context::
@@ -942,6 +943,21 @@ function to temporarily change the active context.
           s = calculate_something()
       s = +s  # Round the final result back to the default precision
 
+   Using keyword arguments, the code would be the following::
+
+      from decimal import localcontext
+
+      with localcontext(prec=42) as ctx:
+          s = calculate_something()
+      s = +s
+
+   Raises :exc:`TypeError` if *kwargs* supplies an attribute that :class:`Context` doesn't
+   support.  Raises either :exc:`TypeError` or :exc:`ValueError` if *kwargs* supplies an
+   invalid value for an attribute.
+
+   .. versionchanged:: 3.11
+      :meth:`localcontext` now supports setting context attributes through the use of keyword arguments.
+
 New contexts can also be created using the :class:`Context` constructor
 described below. In addition, the module provides three pre-made contexts:
 
diff --git a/Lib/_pydecimal.py b/Lib/_pydecimal.py
index 89646fa..f9d6c99 100644
--- a/Lib/_pydecimal.py
+++ b/Lib/_pydecimal.py
@@ -441,6 +441,10 @@ import contextvars
 
 _current_context_var = contextvars.ContextVar('decimal_context')
 
+_context_attributes = frozenset(
+    ['prec', 'Emin', 'Emax', 'capitals', 'clamp', 'rounding', 'flags', 'traps']
+)
+
 def getcontext():
     """Returns this thread's context.
 
@@ -464,7 +468,7 @@ def setcontext(context):
 
 del contextvars        # Don't contaminate the namespace
 
-def localcontext(ctx=None):
+def localcontext(ctx=None, **kwargs):
     """Return a context manager for a copy of the supplied context
 
     Uses a copy of the current context if no context is specified
@@ -500,8 +504,14 @@ def localcontext(ctx=None):
     >>> print(getcontext().prec)
     28
     """
-    if ctx is None: ctx = getcontext()
-    return _ContextManager(ctx)
+    if ctx is None:
+        ctx = getcontext()
+    ctx_manager = _ContextManager(ctx)
+    for key, value in kwargs.items():
+        if key not in _context_attributes:
+            raise TypeError(f"'{key}' is an invalid keyword argument for this function")
+        setattr(ctx_manager.new_context, key, value)
+    return ctx_manager
 
 
 ##### Decimal class #######################################################
diff --git a/Lib/test/test_decimal.py b/Lib/test/test_decimal.py
index 5e77e3c..96f8f7f 100644
--- a/Lib/test/test_decimal.py
+++ b/Lib/test/test_decimal.py
@@ -3665,6 +3665,40 @@ class ContextWithStatement(unittest.TestCase):
         self.assertIsNot(new_ctx, set_ctx, 'did not copy the context')
         self.assertIs(set_ctx, enter_ctx, '__enter__ returned wrong context')
 
+    def test_localcontext_kwargs(self):
+        with self.decimal.localcontext(
+            prec=10, rounding=ROUND_HALF_DOWN,
+            Emin=-20, Emax=20, capitals=0,
+            clamp=1
+        ) as ctx:
+            self.assertEqual(ctx.prec, 10)
+            self.assertEqual(ctx.rounding, self.decimal.ROUND_HALF_DOWN)
+            self.assertEqual(ctx.Emin, -20)
+            self.assertEqual(ctx.Emax, 20)
+            self.assertEqual(ctx.capitals, 0)
+            self.assertEqual(ctx.clamp, 1)
+
+        self.assertRaises(TypeError, self.decimal.localcontext, precision=10)
+
+        self.assertRaises(ValueError, self.decimal.localcontext, Emin=1)
+        self.assertRaises(ValueError, self.decimal.localcontext, Emax=-1)
+        self.assertRaises(ValueError, self.decimal.localcontext, capitals=2)
+        self.assertRaises(ValueError, self.decimal.localcontext, clamp=2)
+
+        self.assertRaises(TypeError, self.decimal.localcontext, rounding="")
+        self.assertRaises(TypeError, self.decimal.localcontext, rounding=1)
+
+        self.assertRaises(TypeError, self.decimal.localcontext, flags="")
+        self.assertRaises(TypeError, self.decimal.localcontext, traps="")
+        self.assertRaises(TypeError, self.decimal.localcontext, Emin="")
+        self.assertRaises(TypeError, self.decimal.localcontext, Emax="")
+
+    def test_local_context_kwargs_does_not_overwrite_existing_argument(self):
+        ctx = self.decimal.getcontext()
+        ctx.prec = 28
+        with self.decimal.localcontext(prec=10) as ctx2:
+            self.assertEqual(ctx.prec, 28)
+
     def test_nested_with_statements(self):
         # Use a copy of the supplied context in the block
         Decimal = self.decimal.Decimal
diff --git a/Misc/NEWS.d/next/Library/2022-04-01-21-44-00.bpo-47135.TvkKB-.rst b/Misc/NEWS.d/next/Library/2022-04-01-21-44-00.bpo-47135.TvkKB-.rst
new file mode 100644
index 0000000..2323c22
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2022-04-01-21-44-00.bpo-47135.TvkKB-.rst
@@ -0,0 +1 @@
+:meth:`decimal.localcontext` now accepts context attributes via keyword arguments
diff --git a/Modules/_decimal/_decimal.c b/Modules/_decimal/_decimal.c
index 4637b8b..8c08847 100644
--- a/Modules/_decimal/_decimal.c
+++ b/Modules/_decimal/_decimal.c
@@ -1156,6 +1156,67 @@ context_setattr(PyObject *self, PyObject *name, PyObject *value)
     return PyObject_GenericSetAttr(self, name, value);
 }
 
+static int
+context_setattrs(PyObject *self, PyObject *prec, PyObject *rounding,
+                 PyObject *emin, PyObject *emax, PyObject *capitals,
+                 PyObject *clamp, PyObject *status, PyObject *traps) {
+
+    int ret;
+    if (prec != Py_None && context_setprec(self, prec, NULL) < 0) {
+        return -1;
+    }
+    if (rounding != Py_None && context_setround(self, rounding, NULL) < 0) {
+        return -1;
+    }
+    if (emin != Py_None && context_setemin(self, emin, NULL) < 0) {
+        return -1;
+    }
+    if (emax != Py_None && context_setemax(self, emax, NULL) < 0) {
+        return -1;
+    }
+    if (capitals != Py_None && context_setcapitals(self, capitals, NULL) < 0) {
+        return -1;
+    }
+    if (clamp != Py_None && context_setclamp(self, clamp, NULL) < 0) {
+       return -1;
+    }
+
+    if (traps != Py_None) {
+        if (PyList_Check(traps)) {
+            ret = context_settraps_list(self, traps);
+        }
+#ifdef EXTRA_FUNCTIONALITY
+        else if (PyLong_Check(traps)) {
+            ret = context_settraps(self, traps, NULL);
+        }
+#endif
+        else {
+            ret = context_settraps_dict(self, traps);
+        }
+        if (ret < 0) {
+            return ret;
+        }
+    }
+    if (status != Py_None) {
+        if (PyList_Check(status)) {
+            ret = context_setstatus_list(self, status);
+        }
+#ifdef EXTRA_FUNCTIONALITY
+        else if (PyLong_Check(status)) {
+            ret = context_setstatus(self, status, NULL);
+        }
+#endif
+        else {
+            ret = context_setstatus_dict(self, status);
+        }
+        if (ret < 0) {
+            return ret;
+        }
+    }
+
+    return 0;
+}
+
 static PyObject *
 context_clear_traps(PyObject *self, PyObject *dummy UNUSED)
 {
@@ -1255,7 +1316,6 @@ context_init(PyObject *self, PyObject *args, PyObject *kwds)
     PyObject *clamp = Py_None;
     PyObject *status = Py_None;
     PyObject *traps = Py_None;
-    int ret;
 
     assert(PyTuple_Check(args));
 
@@ -1267,59 +1327,11 @@ context_init(PyObject *self, PyObject *args, PyObject *kwds)
         return -1;
     }
 
-    if (prec != Py_None && context_setprec(self, prec, NULL) < 0) {
-        return -1;
-    }
-    if (rounding != Py_None && context_setround(self, rounding, NULL) < 0) {
-        return -1;
-    }
-    if (emin != Py_None && context_setemin(self, emin, NULL) < 0) {
-        return -1;
-    }
-    if (emax != Py_None && context_setemax(self, emax, NULL) < 0) {
-        return -1;
-    }
-    if (capitals != Py_None && context_setcapitals(self, capitals, NULL) < 0) {
-        return -1;
-    }
-    if (clamp != Py_None && context_setclamp(self, clamp, NULL) < 0) {
-       return -1;
-    }
-
-    if (traps != Py_None) {
-        if (PyList_Check(traps)) {
-            ret = context_settraps_list(self, traps);
-        }
-#ifdef EXTRA_FUNCTIONALITY
-        else if (PyLong_Check(traps)) {
-            ret = context_settraps(self, traps, NULL);
-        }
-#endif
-        else {
-            ret = context_settraps_dict(self, traps);
-        }
-        if (ret < 0) {
-            return ret;
-        }
-    }
-    if (status != Py_None) {
-        if (PyList_Check(status)) {
-            ret = context_setstatus_list(self, status);
-        }
-#ifdef EXTRA_FUNCTIONALITY
-        else if (PyLong_Check(status)) {
-            ret = context_setstatus(self, status, NULL);
-        }
-#endif
-        else {
-            ret = context_setstatus_dict(self, status);
-        }
-        if (ret < 0) {
-            return ret;
-        }
-    }
-
-    return 0;
+    return context_setattrs(
+        self, prec, rounding,
+        emin, emax, capitals,
+        clamp, status, traps
+    );
 }
 
 static PyObject *
@@ -1721,13 +1733,28 @@ PyDec_SetCurrentContext(PyObject *self UNUSED, PyObject *v)
 static PyObject *
 ctxmanager_new(PyTypeObject *type UNUSED, PyObject *args, PyObject *kwds)
 {
-    static char *kwlist[] = {"ctx", NULL};
+    static char *kwlist[] = {
+      "ctx", "prec", "rounding",
+      "Emin", "Emax", "capitals",
+      "clamp", "flags", "traps",
+      NULL
+    };
     PyDecContextManagerObject *self;
     PyObject *local = Py_None;
     PyObject *global;
 
+    PyObject *prec = Py_None;
+    PyObject *rounding = Py_None;
+    PyObject *Emin = Py_None;
+    PyObject *Emax = Py_None;
+    PyObject *capitals = Py_None;
+    PyObject *clamp = Py_None;
+    PyObject *flags = Py_None;
+    PyObject *traps = Py_None;
+
     CURRENT_CONTEXT(global);
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &local)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOOOOOOOO", kwlist, &local,
+          &prec, &rounding, &Emin, &Emax, &capitals, &clamp, &flags, &traps)) {
         return NULL;
     }
     if (local == Py_None) {
@@ -1754,6 +1781,17 @@ ctxmanager_new(PyTypeObject *type UNUSED, PyObject *args, PyObject *kwds)
     self->global = global;
     Py_INCREF(self->global);
 
+    int ret = context_setattrs(
+        self->local, prec, rounding,
+        Emin, Emax, capitals,
+        clamp, flags, traps
+    );
+
+    if (ret < 0) {
+        Py_DECREF(self);
+        return NULL;
+    }
+
     return (PyObject *)self;
 }
 
diff --git a/Modules/_decimal/docstrings.h b/Modules/_decimal/docstrings.h
index f7fd6e7..a1823cd 100644
--- a/Modules/_decimal/docstrings.h
+++ b/Modules/_decimal/docstrings.h
@@ -30,7 +30,7 @@ Set a new default context.\n\
 \n");
 
 PyDoc_STRVAR(doc_localcontext,
-"localcontext($module, /, ctx=None)\n--\n\n\
+"localcontext($module, /, ctx=None, **kwargs)\n--\n\n\
 Return a context manager that will set the default context to a copy of ctx\n\
 on entry to the with-statement and restore the previous default context when\n\
 exiting the with-statement. If no context is specified, a copy of the current\n\