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

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
This commit is contained in:
Sam Ezeh 2022-04-22 05:27:15 +01:00 committed by GitHub
parent 5e130a8da4
commit bcf14ae433
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 161 additions and 62 deletions

View file

@ -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:

View file

@ -441,6 +441,10 @@ class FloatOperation(DecimalException, TypeError):
_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 sin(x):
>>> 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 #######################################################

View file

@ -3665,6 +3665,40 @@ def test_localcontextarg(self):
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

View file

@ -0,0 +1 @@
:meth:`decimal.localcontext` now accepts context attributes via keyword arguments

View file

@ -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;
}

View file

@ -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\