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 ...@@ -400,7 +400,8 @@ Trigonometric functions
.. function:: dist(p, q) .. function:: dist(p, q)
Return the Euclidean distance between two points *p* and *q*, each 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:: Roughly equivalent to::
......
...@@ -833,6 +833,10 @@ class MathTests(unittest.TestCase): ...@@ -833,6 +833,10 @@ class MathTests(unittest.TestCase):
sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q))) 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__) # Test allowable types (those with __float__)
self.assertEqual(dist((14.0, 1.0), (2.0, -4.0)), 13.0) self.assertEqual(dist((14.0, 1.0), (2.0, -4.0)), 13.0)
self.assertEqual(dist((14, 1), (2, -4)), 13) self.assertEqual(dist((14, 1), (2, -4)), 13)
...@@ -873,8 +877,6 @@ class MathTests(unittest.TestCase): ...@@ -873,8 +877,6 @@ class MathTests(unittest.TestCase):
dist((1, 2, 3), (4, 5, 6), (7, 8, 9)) dist((1, 2, 3), (4, 5, 6), (7, 8, 9))
with self.assertRaises(TypeError): # Scalars not allowed with self.assertRaises(TypeError): # Scalars not allowed
dist(1, 2) 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__ with self.assertRaises(TypeError): # Reject values without __float__
dist((1.1, 'string', 2.2), (1, 2, 3)) dist((1.1, 'string', 2.2), (1, 2, 3))
with self.assertRaises(ValueError): # Check dimension agree 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__, ...@@ -297,8 +297,8 @@ PyDoc_STRVAR(math_dist__doc__,
"\n" "\n"
"Return the Euclidean distance between two points p and q.\n" "Return the Euclidean distance between two points p and q.\n"
"\n" "\n"
"The points should be specified as tuples of coordinates.\n" "The points should be specified as sequences (or iterables) of\n"
"Both tuples must be the same size.\n" "coordinates. Both inputs must have the same dimension.\n"
"\n" "\n"
"Roughly equivalent to:\n" "Roughly equivalent to:\n"
" sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q)))"); " 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) ...@@ -319,15 +319,7 @@ math_dist(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
if (!_PyArg_CheckPositional("dist", nargs, 2, 2)) { if (!_PyArg_CheckPositional("dist", nargs, 2, 2)) {
goto exit; goto exit;
} }
if (!PyTuple_Check(args[0])) {
_PyArg_BadArgument("dist", 1, "tuple", args[0]);
goto exit;
}
p = args[0]; p = args[0];
if (!PyTuple_Check(args[1])) {
_PyArg_BadArgument("dist", 2, "tuple", args[1]);
goto exit;
}
q = args[1]; q = args[1];
return_value = math_dist_impl(module, p, q); return_value = math_dist_impl(module, p, q);
...@@ -720,4 +712,4 @@ math_comb(PyObject *module, PyObject *const *args, Py_ssize_t nargs) ...@@ -720,4 +712,4 @@ math_comb(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
exit: exit:
return return_value; 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) ...@@ -2427,14 +2427,14 @@ vector_norm(Py_ssize_t n, double *vec, double max, int found_nan)
/*[clinic input] /*[clinic input]
math.dist math.dist
p: object(subclass_of='&PyTuple_Type') p: object
q: object(subclass_of='&PyTuple_Type') q: object
/ /
Return the Euclidean distance between two points p and q. Return the Euclidean distance between two points p and q.
The points should be specified as tuples of coordinates. The points should be specified as sequences (or iterables) of
Both tuples must be the same size. coordinates. Both inputs must have the same dimension.
Roughly equivalent to: Roughly equivalent to:
sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q))) sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q)))
...@@ -2442,16 +2442,34 @@ Roughly equivalent to: ...@@ -2442,16 +2442,34 @@ Roughly equivalent to:
static PyObject * static PyObject *
math_dist_impl(PyObject *module, PyObject *p, PyObject *q) 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; PyObject *item;
double max = 0.0; double max = 0.0;
double x, px, qx, result; double x, px, qx, result;
Py_ssize_t i, m, n; 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_on_stack[NUM_STACK_ELEMS];
double *diffs = diffs_on_stack; 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); m = PyTuple_GET_SIZE(p);
n = PyTuple_GET_SIZE(q); n = PyTuple_GET_SIZE(q);
if (m != n) { if (m != n) {
...@@ -2482,12 +2500,24 @@ math_dist_impl(PyObject *module, PyObject *p, PyObject *q) ...@@ -2482,12 +2500,24 @@ math_dist_impl(PyObject *module, PyObject *p, PyObject *q)
if (diffs != diffs_on_stack) { if (diffs != diffs_on_stack) {
PyObject_Free(diffs); PyObject_Free(diffs);
} }
if (p_allocated) {
Py_DECREF(p);
}
if (q_allocated) {
Py_DECREF(q);
}
return PyFloat_FromDouble(result); return PyFloat_FromDouble(result);
error_exit: error_exit:
if (diffs != diffs_on_stack) { if (diffs != diffs_on_stack) {
PyObject_Free(diffs); PyObject_Free(diffs);
} }
if (p_allocated) {
Py_DECREF(p);
}
if (q_allocated) {
Py_DECREF(q);
}
return NULL; 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