bpo-35606: Implement math.prod (GH-11359)

This commit is contained in:
Pablo Galindo 2019-02-07 07:04:02 +00:00 committed by Raymond Hettinger
parent e9bc4172d1
commit bc09851586
6 changed files with 260 additions and 1 deletions

View File

@ -178,6 +178,18 @@ Number-theoretic and representation functions
of *x* and are floats.
.. function:: prod(iterable, *, start=1)
Calculate the product of all the elements in the input *iterable*.
The default *start* value for the product is ``1``.
When the iterable is empty, return the start value. This function is
intended specifically for use with numeric values and may reject
non-numeric types.
.. versionadded:: 3.8
.. function:: remainder(x, y)
Return the IEEE 754-style remainder of *x* with respect to *y*. For

View File

@ -171,6 +171,15 @@ json.tool
Add option ``--json-lines`` to parse every input line as separate JSON object.
(Contributed by Weipeng Hong in :issue:`31553`.)
math
----
Added new function, :func:`math.prod`, as analogous function to :func:`sum`
that returns the product of a 'start' value (default: 1) times an iterable of
numbers. (Contributed by Pablo Galindo in :issue:`issue35606`)
os.path
-------

View File

@ -1724,6 +1724,37 @@ class IsCloseTests(unittest.TestCase):
self.assertAllClose(fraction_examples, rel_tol=1e-8)
self.assertAllNotClose(fraction_examples, rel_tol=1e-9)
def test_prod(self):
prod = math.prod
self.assertEqual(prod([]), 1)
self.assertEqual(prod([], start=5), 5)
self.assertEqual(prod(list(range(2,8))), 5040)
self.assertEqual(prod(iter(list(range(2,8)))), 5040)
self.assertEqual(prod(range(1, 10), start=10), 3628800)
self.assertEqual(prod([1, 2, 3, 4, 5]), 120)
self.assertEqual(prod([1.0, 2.0, 3.0, 4.0, 5.0]), 120.0)
self.assertEqual(prod([1, 2, 3, 4.0, 5.0]), 120.0)
self.assertEqual(prod([1.0, 2.0, 3.0, 4, 5]), 120.0)
# Test overflow in fast-path for integers
self.assertEqual(prod([1, 1, 2**32, 1, 1]), 2**32)
# Test overflow in fast-path for floats
self.assertEqual(prod([1.0, 1.0, 2**32, 1, 1]), float(2**32))
self.assertRaises(TypeError, prod)
self.assertRaises(TypeError, prod, 42)
self.assertRaises(TypeError, prod, ['a', 'b', 'c'])
self.assertRaises(TypeError, prod, ['a', 'b', 'c'], '')
self.assertRaises(TypeError, prod, [b'a', b'c'], b'')
values = [bytearray(b'a'), bytearray(b'b')]
self.assertRaises(TypeError, prod, values, bytearray(b''))
self.assertRaises(TypeError, prod, [[1], [2], [3]])
self.assertRaises(TypeError, prod, [{2:3}])
self.assertRaises(TypeError, prod, [{2:3}]*2, {2:3})
self.assertRaises(TypeError, prod, [[1], [2], [3]], [])
with self.assertRaises(TypeError):
prod([10, 20], [30, 40]) # start is a keyword-only argument
def test_main():
from doctest import DocFileSuite

View File

@ -0,0 +1,3 @@
Implement :func:`math.prod` as analogous function to :func:`sum` that
returns the product of a 'start' value (default: 1) times an iterable of
numbers. Patch by Pablo Galindo.

View File

@ -556,4 +556,41 @@ math_isclose(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject
exit:
return return_value;
}
/*[clinic end generated code: output=0664f30046da09fe input=a9049054013a1b77]*/
PyDoc_STRVAR(math_prod__doc__,
"prod($module, iterable, /, *, start=1)\n"
"--\n"
"\n"
"Calculate the product of all the elements in the input iterable.\n"
"\n"
"The default start value for the product is 1.\n"
"\n"
"When the iterable is empty, return the start value. This function is\n"
"intended specifically for use with numeric values and may reject\n"
"non-numeric types.");
#define MATH_PROD_METHODDEF \
{"prod", (PyCFunction)(void(*)(void))math_prod, METH_FASTCALL|METH_KEYWORDS, math_prod__doc__},
static PyObject *
math_prod_impl(PyObject *module, PyObject *iterable, PyObject *start);
static PyObject *
math_prod(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
{
PyObject *return_value = NULL;
static const char * const _keywords[] = {"", "start", NULL};
static _PyArg_Parser _parser = {"O|$O:prod", _keywords, 0};
PyObject *iterable;
PyObject *start = NULL;
if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
&iterable, &start)) {
goto exit;
}
return_value = math_prod_impl(module, iterable, start);
exit:
return return_value;
}
/*[clinic end generated code: output=20505690ca6fe402 input=a9049054013a1b77]*/

View File

@ -2494,6 +2494,172 @@ math_isclose_impl(PyObject *module, double a, double b, double rel_tol,
}
/*[clinic input]
math.prod
iterable: object
/
*
start: object(c_default="NULL") = 1
Calculate the product of all the elements in the input iterable.
The default start value for the product is 1.
When the iterable is empty, return the start value. This function is
intended specifically for use with numeric values and may reject
non-numeric types.
[clinic start generated code]*/
static PyObject *
math_prod_impl(PyObject *module, PyObject *iterable, PyObject *start)
/*[clinic end generated code: output=36153bedac74a198 input=4c5ab0682782ed54]*/
{
PyObject *result = start;
PyObject *temp, *item, *iter;
iter = PyObject_GetIter(iterable);
if (iter == NULL) {
return NULL;
}
if (result == NULL) {
result = PyLong_FromLong(1);
if (result == NULL) {
Py_DECREF(iter);
return NULL;
}
} else {
Py_INCREF(result);
}
#ifndef SLOW_PROD
/* Fast paths for integers keeping temporary products in C.
* Assumes all inputs are the same type.
* If the assumption fails, default to use PyObjects instead.
*/
if (PyLong_CheckExact(result)) {
int overflow;
long i_result = PyLong_AsLongAndOverflow(result, &overflow);
/* If this already overflowed, don't even enter the loop. */
if (overflow == 0) {
Py_DECREF(result);
result = NULL;
}
/* Loop over all the items in the iterable until we finish, we overflow
* or we found a non integer element */
while(result == NULL) {
item = PyIter_Next(iter);
if (item == NULL) {
Py_DECREF(iter);
if (PyErr_Occurred()) {
return NULL;
}
return PyLong_FromLong(i_result);
}
if (PyLong_CheckExact(item)) {
long b = PyLong_AsLongAndOverflow(item, &overflow);
long x = i_result * b;
/* Continue if there is no overflow */
if (overflow == 0
&& x < INT_MAX && x > INT_MIN
&& !(b != 0 && x / i_result != b)) {
i_result = x;
Py_DECREF(item);
continue;
}
}
/* Either overflowed or is not an int.
* Restore real objects and process normally */
result = PyLong_FromLong(i_result);
if (result == NULL) {
Py_DECREF(item);
Py_DECREF(iter);
return NULL;
}
temp = PyNumber_Multiply(result, item);
Py_DECREF(result);
Py_DECREF(item);
result = temp;
if (result == NULL) {
Py_DECREF(iter);
return NULL;
}
}
}
/* Fast paths for floats keeping temporary products in C.
* Assumes all inputs are the same type.
* If the assumption fails, default to use PyObjects instead.
*/
if (PyFloat_CheckExact(result)) {
double f_result = PyFloat_AS_DOUBLE(result);
Py_DECREF(result);
result = NULL;
while(result == NULL) {
item = PyIter_Next(iter);
if (item == NULL) {
Py_DECREF(iter);
if (PyErr_Occurred()) {
return NULL;
}
return PyFloat_FromDouble(f_result);
}
if (PyFloat_CheckExact(item)) {
f_result *= PyFloat_AS_DOUBLE(item);
Py_DECREF(item);
continue;
}
if (PyLong_CheckExact(item)) {
long value;
int overflow;
value = PyLong_AsLongAndOverflow(item, &overflow);
if (!overflow) {
f_result *= (double)value;
Py_DECREF(item);
continue;
}
}
result = PyFloat_FromDouble(f_result);
if (result == NULL) {
Py_DECREF(item);
Py_DECREF(iter);
return NULL;
}
temp = PyNumber_Multiply(result, item);
Py_DECREF(result);
Py_DECREF(item);
result = temp;
if (result == NULL) {
Py_DECREF(iter);
return NULL;
}
}
}
#endif
/* Consume rest of the iterable (if any) that could not be handled
* by specialized functions above.*/
for(;;) {
item = PyIter_Next(iter);
if (item == NULL) {
/* error, or end-of-sequence */
if (PyErr_Occurred()) {
Py_DECREF(result);
result = NULL;
}
break;
}
temp = PyNumber_Multiply(result, item);
Py_DECREF(result);
Py_DECREF(item);
result = temp;
if (result == NULL)
break;
}
Py_DECREF(iter);
return result;
}
static PyMethodDef math_methods[] = {
{"acos", math_acos, METH_O, math_acos_doc},
{"acosh", math_acosh, METH_O, math_acosh_doc},
@ -2541,6 +2707,7 @@ static PyMethodDef math_methods[] = {
{"tan", math_tan, METH_O, math_tan_doc},
{"tanh", math_tanh, METH_O, math_tanh_doc},
MATH_TRUNC_METHODDEF
MATH_PROD_METHODDEF
{NULL, NULL} /* sentinel */
};