Commit 6b5f1b49 authored by Raymond Hettinger's avatar Raymond Hettinger Committed by GitHub

bpo-37691: Let math.dist() accept sequences and iterables for coordinates (GH-14975)

parent 3221a63c
......@@ -400,7 +400,8 @@ Trigonometric functions
.. function:: dist(p, q)
Return the Euclidean distance between two points *p* and *q*, each
given as a tuple of coordinates. The two tuples must be the same size.
given as a sequence (or iterable) of coordinates. The two points
must have the same dimension.
Roughly equivalent to::
......
......@@ -833,6 +833,10 @@ class MathTests(unittest.TestCase):
sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q)))
)
# Test non-tuple inputs
self.assertEqual(dist([1.0, 2.0, 3.0], [4.0, 2.0, -1.0]), 5.0)
self.assertEqual(dist(iter([1.0, 2.0, 3.0]), iter([4.0, 2.0, -1.0])), 5.0)
# Test allowable types (those with __float__)
self.assertEqual(dist((14.0, 1.0), (2.0, -4.0)), 13.0)
self.assertEqual(dist((14, 1), (2, -4)), 13)
......@@ -873,8 +877,6 @@ class MathTests(unittest.TestCase):
dist((1, 2, 3), (4, 5, 6), (7, 8, 9))
with self.assertRaises(TypeError): # Scalars not allowed
dist(1, 2)
with self.assertRaises(TypeError): # Lists not allowed
dist([1, 2, 3], [4, 5, 6])
with self.assertRaises(TypeError): # Reject values without __float__
dist((1.1, 'string', 2.2), (1, 2, 3))
with self.assertRaises(ValueError): # Check dimension agree
......
Let math.dist() accept coordinates as sequences (or iterables) rather than
just tuples.
......@@ -297,8 +297,8 @@ PyDoc_STRVAR(math_dist__doc__,
"\n"
"Return the Euclidean distance between two points p and q.\n"
"\n"
"The points should be specified as tuples of coordinates.\n"
"Both tuples must be the same size.\n"
"The points should be specified as sequences (or iterables) of\n"
"coordinates. Both inputs must have the same dimension.\n"
"\n"
"Roughly equivalent to:\n"
" sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q)))");
......@@ -319,15 +319,7 @@ math_dist(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
if (!_PyArg_CheckPositional("dist", nargs, 2, 2)) {
goto exit;
}
if (!PyTuple_Check(args[0])) {
_PyArg_BadArgument("dist", 1, "tuple", args[0]);
goto exit;
}
p = args[0];
if (!PyTuple_Check(args[1])) {
_PyArg_BadArgument("dist", 2, "tuple", args[1]);
goto exit;
}
q = args[1];
return_value = math_dist_impl(module, p, q);
......@@ -720,4 +712,4 @@ math_comb(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
exit:
return return_value;
}
/*[clinic end generated code: output=0eb1e76a769cdd30 input=a9049054013a1b77]*/
/*[clinic end generated code: output=f93cfe13ab2fdb4e input=a9049054013a1b77]*/
......@@ -2427,14 +2427,14 @@ vector_norm(Py_ssize_t n, double *vec, double max, int found_nan)
/*[clinic input]
math.dist
p: object(subclass_of='&PyTuple_Type')
q: object(subclass_of='&PyTuple_Type')
p: object
q: object
/
Return the Euclidean distance between two points p and q.
The points should be specified as tuples of coordinates.
Both tuples must be the same size.
The points should be specified as sequences (or iterables) of
coordinates. Both inputs must have the same dimension.
Roughly equivalent to:
sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q)))
......@@ -2442,16 +2442,34 @@ Roughly equivalent to:
static PyObject *
math_dist_impl(PyObject *module, PyObject *p, PyObject *q)
/*[clinic end generated code: output=56bd9538d06bbcfe input=937122eaa5f19272]*/
/*[clinic end generated code: output=56bd9538d06bbcfe input=74e85e1b6092e68e]*/
{
PyObject *item;
double max = 0.0;
double x, px, qx, result;
Py_ssize_t i, m, n;
int found_nan = 0;
int found_nan = 0, p_allocated = 0, q_allocated = 0;
double diffs_on_stack[NUM_STACK_ELEMS];
double *diffs = diffs_on_stack;
if (!PyTuple_Check(p)) {
p = PySequence_Tuple(p);
if (p == NULL) {
return NULL;
}
p_allocated = 1;
}
if (!PyTuple_Check(q)) {
q = PySequence_Tuple(q);
if (q == NULL) {
if (p_allocated) {
Py_DECREF(p);
}
return NULL;
}
q_allocated = 1;
}
m = PyTuple_GET_SIZE(p);
n = PyTuple_GET_SIZE(q);
if (m != n) {
......@@ -2482,12 +2500,24 @@ math_dist_impl(PyObject *module, PyObject *p, PyObject *q)
if (diffs != diffs_on_stack) {
PyObject_Free(diffs);
}
if (p_allocated) {
Py_DECREF(p);
}
if (q_allocated) {
Py_DECREF(q);
}
return PyFloat_FromDouble(result);
error_exit:
if (diffs != diffs_on_stack) {
PyObject_Free(diffs);
}
if (p_allocated) {
Py_DECREF(p);
}
if (q_allocated) {
Py_DECREF(q);
}
return NULL;
}
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment