Skip to content

Commit bcf14ae

Browse files
gh-91291: Accept attributes as keyword arguments in decimal.localcontext (#32242)
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
1 parent 5e130a8 commit bcf14ae

File tree

6 files changed

+161
-62
lines changed

6 files changed

+161
-62
lines changed

Doc/library/decimal.rst

+18-2
Original file line numberDiff line numberDiff line change
@@ -925,12 +925,13 @@ Each thread has its own current context which is accessed or changed using the
925925
You can also use the :keyword:`with` statement and the :func:`localcontext`
926926
function to temporarily change the active context.
927927

928-
.. function:: localcontext(ctx=None)
928+
.. function:: localcontext(ctx=None, \*\*kwargs)
929929

930930
Return a context manager that will set the current context for the active thread
931931
to a copy of *ctx* on entry to the with-statement and restore the previous context
932932
when exiting the with-statement. If no context is specified, a copy of the
933-
current context is used.
933+
current context is used. The *kwargs* argument is used to set the attributes
934+
of the new context.
934935

935936
For example, the following code sets the current decimal precision to 42 places,
936937
performs a calculation, and then automatically restores the previous context::
@@ -942,6 +943,21 @@ function to temporarily change the active context.
942943
s = calculate_something()
943944
s = +s # Round the final result back to the default precision
944945

946+
Using keyword arguments, the code would be the following::
947+
948+
from decimal import localcontext
949+
950+
with localcontext(prec=42) as ctx:
951+
s = calculate_something()
952+
s = +s
953+
954+
Raises :exc:`TypeError` if *kwargs* supplies an attribute that :class:`Context` doesn't
955+
support. Raises either :exc:`TypeError` or :exc:`ValueError` if *kwargs* supplies an
956+
invalid value for an attribute.
957+
958+
.. versionchanged:: 3.11
959+
:meth:`localcontext` now supports setting context attributes through the use of keyword arguments.
960+
945961
New contexts can also be created using the :class:`Context` constructor
946962
described below. In addition, the module provides three pre-made contexts:
947963

Lib/_pydecimal.py

+13-3
Original file line numberDiff line numberDiff line change
@@ -441,6 +441,10 @@ class FloatOperation(DecimalException, TypeError):
441441

442442
_current_context_var = contextvars.ContextVar('decimal_context')
443443

444+
_context_attributes = frozenset(
445+
['prec', 'Emin', 'Emax', 'capitals', 'clamp', 'rounding', 'flags', 'traps']
446+
)
447+
444448
def getcontext():
445449
"""Returns this thread's context.
446450
@@ -464,7 +468,7 @@ def setcontext(context):
464468

465469
del contextvars # Don't contaminate the namespace
466470

467-
def localcontext(ctx=None):
471+
def localcontext(ctx=None, **kwargs):
468472
"""Return a context manager for a copy of the supplied context
469473
470474
Uses a copy of the current context if no context is specified
@@ -500,8 +504,14 @@ def sin(x):
500504
>>> print(getcontext().prec)
501505
28
502506
"""
503-
if ctx is None: ctx = getcontext()
504-
return _ContextManager(ctx)
507+
if ctx is None:
508+
ctx = getcontext()
509+
ctx_manager = _ContextManager(ctx)
510+
for key, value in kwargs.items():
511+
if key not in _context_attributes:
512+
raise TypeError(f"'{key}' is an invalid keyword argument for this function")
513+
setattr(ctx_manager.new_context, key, value)
514+
return ctx_manager
505515

506516

507517
##### Decimal class #######################################################

Lib/test/test_decimal.py

+34
Original file line numberDiff line numberDiff line change
@@ -3665,6 +3665,40 @@ def test_localcontextarg(self):
36653665
self.assertIsNot(new_ctx, set_ctx, 'did not copy the context')
36663666
self.assertIs(set_ctx, enter_ctx, '__enter__ returned wrong context')
36673667

3668+
def test_localcontext_kwargs(self):
3669+
with self.decimal.localcontext(
3670+
prec=10, rounding=ROUND_HALF_DOWN,
3671+
Emin=-20, Emax=20, capitals=0,
3672+
clamp=1
3673+
) as ctx:
3674+
self.assertEqual(ctx.prec, 10)
3675+
self.assertEqual(ctx.rounding, self.decimal.ROUND_HALF_DOWN)
3676+
self.assertEqual(ctx.Emin, -20)
3677+
self.assertEqual(ctx.Emax, 20)
3678+
self.assertEqual(ctx.capitals, 0)
3679+
self.assertEqual(ctx.clamp, 1)
3680+
3681+
self.assertRaises(TypeError, self.decimal.localcontext, precision=10)
3682+
3683+
self.assertRaises(ValueError, self.decimal.localcontext, Emin=1)
3684+
self.assertRaises(ValueError, self.decimal.localcontext, Emax=-1)
3685+
self.assertRaises(ValueError, self.decimal.localcontext, capitals=2)
3686+
self.assertRaises(ValueError, self.decimal.localcontext, clamp=2)
3687+
3688+
self.assertRaises(TypeError, self.decimal.localcontext, rounding="")
3689+
self.assertRaises(TypeError, self.decimal.localcontext, rounding=1)
3690+
3691+
self.assertRaises(TypeError, self.decimal.localcontext, flags="")
3692+
self.assertRaises(TypeError, self.decimal.localcontext, traps="")
3693+
self.assertRaises(TypeError, self.decimal.localcontext, Emin="")
3694+
self.assertRaises(TypeError, self.decimal.localcontext, Emax="")
3695+
3696+
def test_local_context_kwargs_does_not_overwrite_existing_argument(self):
3697+
ctx = self.decimal.getcontext()
3698+
ctx.prec = 28
3699+
with self.decimal.localcontext(prec=10) as ctx2:
3700+
self.assertEqual(ctx.prec, 28)
3701+
36683702
def test_nested_with_statements(self):
36693703
# Use a copy of the supplied context in the block
36703704
Decimal = self.decimal.Decimal
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
:meth:`decimal.localcontext` now accepts context attributes via keyword arguments

Modules/_decimal/_decimal.c

+94-56
Original file line numberDiff line numberDiff line change
@@ -1156,6 +1156,67 @@ context_setattr(PyObject *self, PyObject *name, PyObject *value)
11561156
return PyObject_GenericSetAttr(self, name, value);
11571157
}
11581158

1159+
static int
1160+
context_setattrs(PyObject *self, PyObject *prec, PyObject *rounding,
1161+
PyObject *emin, PyObject *emax, PyObject *capitals,
1162+
PyObject *clamp, PyObject *status, PyObject *traps) {
1163+
1164+
int ret;
1165+
if (prec != Py_None && context_setprec(self, prec, NULL) < 0) {
1166+
return -1;
1167+
}
1168+
if (rounding != Py_None && context_setround(self, rounding, NULL) < 0) {
1169+
return -1;
1170+
}
1171+
if (emin != Py_None && context_setemin(self, emin, NULL) < 0) {
1172+
return -1;
1173+
}
1174+
if (emax != Py_None && context_setemax(self, emax, NULL) < 0) {
1175+
return -1;
1176+
}
1177+
if (capitals != Py_None && context_setcapitals(self, capitals, NULL) < 0) {
1178+
return -1;
1179+
}
1180+
if (clamp != Py_None && context_setclamp(self, clamp, NULL) < 0) {
1181+
return -1;
1182+
}
1183+
1184+
if (traps != Py_None) {
1185+
if (PyList_Check(traps)) {
1186+
ret = context_settraps_list(self, traps);
1187+
}
1188+
#ifdef EXTRA_FUNCTIONALITY
1189+
else if (PyLong_Check(traps)) {
1190+
ret = context_settraps(self, traps, NULL);
1191+
}
1192+
#endif
1193+
else {
1194+
ret = context_settraps_dict(self, traps);
1195+
}
1196+
if (ret < 0) {
1197+
return ret;
1198+
}
1199+
}
1200+
if (status != Py_None) {
1201+
if (PyList_Check(status)) {
1202+
ret = context_setstatus_list(self, status);
1203+
}
1204+
#ifdef EXTRA_FUNCTIONALITY
1205+
else if (PyLong_Check(status)) {
1206+
ret = context_setstatus(self, status, NULL);
1207+
}
1208+
#endif
1209+
else {
1210+
ret = context_setstatus_dict(self, status);
1211+
}
1212+
if (ret < 0) {
1213+
return ret;
1214+
}
1215+
}
1216+
1217+
return 0;
1218+
}
1219+
11591220
static PyObject *
11601221
context_clear_traps(PyObject *self, PyObject *dummy UNUSED)
11611222
{
@@ -1255,7 +1316,6 @@ context_init(PyObject *self, PyObject *args, PyObject *kwds)
12551316
PyObject *clamp = Py_None;
12561317
PyObject *status = Py_None;
12571318
PyObject *traps = Py_None;
1258-
int ret;
12591319

12601320
assert(PyTuple_Check(args));
12611321

@@ -1267,59 +1327,11 @@ context_init(PyObject *self, PyObject *args, PyObject *kwds)
12671327
return -1;
12681328
}
12691329

1270-
if (prec != Py_None && context_setprec(self, prec, NULL) < 0) {
1271-
return -1;
1272-
}
1273-
if (rounding != Py_None && context_setround(self, rounding, NULL) < 0) {
1274-
return -1;
1275-
}
1276-
if (emin != Py_None && context_setemin(self, emin, NULL) < 0) {
1277-
return -1;
1278-
}
1279-
if (emax != Py_None && context_setemax(self, emax, NULL) < 0) {
1280-
return -1;
1281-
}
1282-
if (capitals != Py_None && context_setcapitals(self, capitals, NULL) < 0) {
1283-
return -1;
1284-
}
1285-
if (clamp != Py_None && context_setclamp(self, clamp, NULL) < 0) {
1286-
return -1;
1287-
}
1288-
1289-
if (traps != Py_None) {
1290-
if (PyList_Check(traps)) {
1291-
ret = context_settraps_list(self, traps);
1292-
}
1293-
#ifdef EXTRA_FUNCTIONALITY
1294-
else if (PyLong_Check(traps)) {
1295-
ret = context_settraps(self, traps, NULL);
1296-
}
1297-
#endif
1298-
else {
1299-
ret = context_settraps_dict(self, traps);
1300-
}
1301-
if (ret < 0) {
1302-
return ret;
1303-
}
1304-
}
1305-
if (status != Py_None) {
1306-
if (PyList_Check(status)) {
1307-
ret = context_setstatus_list(self, status);
1308-
}
1309-
#ifdef EXTRA_FUNCTIONALITY
1310-
else if (PyLong_Check(status)) {
1311-
ret = context_setstatus(self, status, NULL);
1312-
}
1313-
#endif
1314-
else {
1315-
ret = context_setstatus_dict(self, status);
1316-
}
1317-
if (ret < 0) {
1318-
return ret;
1319-
}
1320-
}
1321-
1322-
return 0;
1330+
return context_setattrs(
1331+
self, prec, rounding,
1332+
emin, emax, capitals,
1333+
clamp, status, traps
1334+
);
13231335
}
13241336

13251337
static PyObject *
@@ -1721,13 +1733,28 @@ PyDec_SetCurrentContext(PyObject *self UNUSED, PyObject *v)
17211733
static PyObject *
17221734
ctxmanager_new(PyTypeObject *type UNUSED, PyObject *args, PyObject *kwds)
17231735
{
1724-
static char *kwlist[] = {"ctx", NULL};
1736+
static char *kwlist[] = {
1737+
"ctx", "prec", "rounding",
1738+
"Emin", "Emax", "capitals",
1739+
"clamp", "flags", "traps",
1740+
NULL
1741+
};
17251742
PyDecContextManagerObject *self;
17261743
PyObject *local = Py_None;
17271744
PyObject *global;
17281745

1746+
PyObject *prec = Py_None;
1747+
PyObject *rounding = Py_None;
1748+
PyObject *Emin = Py_None;
1749+
PyObject *Emax = Py_None;
1750+
PyObject *capitals = Py_None;
1751+
PyObject *clamp = Py_None;
1752+
PyObject *flags = Py_None;
1753+
PyObject *traps = Py_None;
1754+
17291755
CURRENT_CONTEXT(global);
1730-
if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &local)) {
1756+
if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOOOOOOOO", kwlist, &local,
1757+
&prec, &rounding, &Emin, &Emax, &capitals, &clamp, &flags, &traps)) {
17311758
return NULL;
17321759
}
17331760
if (local == Py_None) {
@@ -1754,6 +1781,17 @@ ctxmanager_new(PyTypeObject *type UNUSED, PyObject *args, PyObject *kwds)
17541781
self->global = global;
17551782
Py_INCREF(self->global);
17561783

1784+
int ret = context_setattrs(
1785+
self->local, prec, rounding,
1786+
Emin, Emax, capitals,
1787+
clamp, flags, traps
1788+
);
1789+
1790+
if (ret < 0) {
1791+
Py_DECREF(self);
1792+
return NULL;
1793+
}
1794+
17571795
return (PyObject *)self;
17581796
}
17591797

Modules/_decimal/docstrings.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ Set a new default context.\n\
3030
\n");
3131

3232
PyDoc_STRVAR(doc_localcontext,
33-
"localcontext($module, /, ctx=None)\n--\n\n\
33+
"localcontext($module, /, ctx=None, **kwargs)\n--\n\n\
3434
Return a context manager that will set the default context to a copy of ctx\n\
3535
on entry to the with-statement and restore the previous default context when\n\
3636
exiting the with-statement. If no context is specified, a copy of the current\n\

0 commit comments

Comments
 (0)