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` You can also use the :keyword:`with` statement and the :func:`localcontext`
function to temporarily change the active context. 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 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 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 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, For example, the following code sets the current decimal precision to 42 places,
performs a calculation, and then automatically restores the previous context:: 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 = calculate_something()
s = +s # Round the final result back to the default precision 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 New contexts can also be created using the :class:`Context` constructor
described below. In addition, the module provides three pre-made contexts: described below. In addition, the module provides three pre-made contexts:

View File

@ -441,6 +441,10 @@ import contextvars
_current_context_var = contextvars.ContextVar('decimal_context') _current_context_var = contextvars.ContextVar('decimal_context')
_context_attributes = frozenset(
['prec', 'Emin', 'Emax', 'capitals', 'clamp', 'rounding', 'flags', 'traps']
)
def getcontext(): def getcontext():
"""Returns this thread's context. """Returns this thread's context.
@ -464,7 +468,7 @@ def setcontext(context):
del contextvars # Don't contaminate the namespace 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 """Return a context manager for a copy of the supplied context
Uses a copy of the current context if no context is specified Uses a copy of the current context if no context is specified
@ -500,8 +504,14 @@ def localcontext(ctx=None):
>>> print(getcontext().prec) >>> print(getcontext().prec)
28 28
""" """
if ctx is None: ctx = getcontext() if ctx is None:
return _ContextManager(ctx) 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 ####################################################### ##### Decimal class #######################################################

View File

@ -3665,6 +3665,40 @@ class ContextWithStatement(unittest.TestCase):
self.assertIsNot(new_ctx, set_ctx, 'did not copy the context') self.assertIsNot(new_ctx, set_ctx, 'did not copy the context')
self.assertIs(set_ctx, enter_ctx, '__enter__ returned wrong 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): def test_nested_with_statements(self):
# Use a copy of the supplied context in the block # Use a copy of the supplied context in the block
Decimal = self.decimal.Decimal 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); 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 * static PyObject *
context_clear_traps(PyObject *self, PyObject *dummy UNUSED) 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 *clamp = Py_None;
PyObject *status = Py_None; PyObject *status = Py_None;
PyObject *traps = Py_None; PyObject *traps = Py_None;
int ret;
assert(PyTuple_Check(args)); assert(PyTuple_Check(args));
@ -1267,59 +1327,11 @@ context_init(PyObject *self, PyObject *args, PyObject *kwds)
return -1; return -1;
} }
if (prec != Py_None && context_setprec(self, prec, NULL) < 0) { return context_setattrs(
return -1; self, prec, rounding,
} emin, emax, capitals,
if (rounding != Py_None && context_setround(self, rounding, NULL) < 0) { clamp, status, traps
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 * static PyObject *
@ -1721,13 +1733,28 @@ PyDec_SetCurrentContext(PyObject *self UNUSED, PyObject *v)
static PyObject * static PyObject *
ctxmanager_new(PyTypeObject *type UNUSED, PyObject *args, PyObject *kwds) 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; PyDecContextManagerObject *self;
PyObject *local = Py_None; PyObject *local = Py_None;
PyObject *global; 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); 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; return NULL;
} }
if (local == Py_None) { if (local == Py_None) {
@ -1754,6 +1781,17 @@ ctxmanager_new(PyTypeObject *type UNUSED, PyObject *args, PyObject *kwds)
self->global = global; self->global = global;
Py_INCREF(self->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; return (PyObject *)self;
} }

View File

@ -30,7 +30,7 @@ Set a new default context.\n\
\n"); \n");
PyDoc_STRVAR(doc_localcontext, 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\ 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\ 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\ exiting the with-statement. If no context is specified, a copy of the current\n\