bpo-33089: Add math.dist() for computing the Euclidean distance between two points (GH-8561)

2 files changed, 120 insertions, 1 deletions

diff --git a/Modules/clinic/mathmodule.c.h b/Modules/clinic/mathmodule.c.h
index 3a487bd..c4d2786 100644
--- a/Modules/clinic/mathmodule.c.h
+++ b/Modules/clinic/mathmodule.c.h
@@ -269,6 +269,41 @@ exit:
     return return_value;
 }
 
+PyDoc_STRVAR(math_dist__doc__,
+"dist($module, p, q, /)\n"
+"--\n"
+"\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"
+"\n"
+"Roughly equivalent to:\n"
+"    sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q)))");
+
+#define MATH_DIST_METHODDEF    \
+    {"dist", (PyCFunction)math_dist, METH_FASTCALL, math_dist__doc__},
+
+static PyObject *
+math_dist_impl(PyObject *module, PyObject *p, PyObject *q);
+
+static PyObject *
+math_dist(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    PyObject *p;
+    PyObject *q;
+
+    if (!_PyArg_ParseStack(args, nargs, "O!O!:dist",
+        &PyTuple_Type, &p, &PyTuple_Type, &q)) {
+        goto exit;
+    }
+    return_value = math_dist_impl(module, p, q);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(math_pow__doc__,
 "pow($module, x, y, /)\n"
 "--\n"
@@ -487,4 +522,4 @@ math_isclose(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=1c35516a10443902 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=d936137c1189b89b input=a9049054013a1b77]*/
diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c
index 726fc56..8f11f2c 100644
--- a/Modules/mathmodule.c
+++ b/Modules/mathmodule.c
@@ -2031,6 +2031,89 @@ math_fmod_impl(PyObject *module, double x, double y)
         return PyFloat_FromDouble(r);
 }
 
+/*[clinic input]
+math.dist
+
+    p: object(subclass_of='&PyTuple_Type')
+    q: object(subclass_of='&PyTuple_Type')
+    /
+
+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.
+
+Roughly equivalent to:
+    sqrt(sum((px - qx) ** 2.0 for px, qx in zip(p, q)))
+[clinic start generated code]*/
+
+static PyObject *
+math_dist_impl(PyObject *module, PyObject *p, PyObject *q)
+/*[clinic end generated code: output=56bd9538d06bbcfe input=937122eaa5f19272]*/
+{
+    PyObject *item;
+    double *diffs;
+    double max = 0.0;
+    double csum = 0.0;
+    double x, px, qx, result;
+    Py_ssize_t i, m, n;
+    int found_nan = 0;
+
+    m = PyTuple_GET_SIZE(p);
+    n = PyTuple_GET_SIZE(q);
+    if (m != n) {
+        PyErr_SetString(PyExc_ValueError,
+                        "both points must have the same number of dimensions");
+        return NULL;
+
+    }
+    diffs = (double *) PyObject_Malloc(n * sizeof(double));
+    if (diffs == NULL) {
+        return NULL;
+    }
+    for (i=0 ; i<n ; i++) {
+        item = PyTuple_GET_ITEM(p, i);
+        px = PyFloat_AsDouble(item);
+        if (px == -1.0 && PyErr_Occurred()) {
+            PyObject_Free(diffs);
+            return NULL;
+        }
+        item = PyTuple_GET_ITEM(q, i);
+        qx = PyFloat_AsDouble(item);
+        if (qx == -1.0 && PyErr_Occurred()) {
+            PyObject_Free(diffs);
+            return NULL;
+        }
+        x = fabs(px - qx);
+        diffs[i] = x;
+        found_nan |= Py_IS_NAN(x);
+        if (x > max) {
+            max = x;
+        }
+    }
+    if (Py_IS_INFINITY(max)) {
+        result = max;
+        goto done;
+    }
+    if (found_nan) {
+        result = Py_NAN;
+        goto done;
+    }
+    if (max == 0.0) {
+        result = 0.0;
+        goto done;
+    }
+    for (i=0 ; i<n ; i++) {
+        x = diffs[i] / max;
+        csum += x * x;
+    }
+    result = max * sqrt(csum);
+
+  done:
+    PyObject_Free(diffs);
+    return PyFloat_FromDouble(result);
+}
+
 /* AC: cannot convert yet, waiting for *args support */
 static PyObject *
 math_hypot(PyObject *self, PyObject *args)
@@ -2358,6 +2441,7 @@ static PyMethodDef math_methods[] = {
     {"cos",             math_cos,       METH_O,         math_cos_doc},
     {"cosh",            math_cosh,      METH_O,         math_cosh_doc},
     MATH_DEGREES_METHODDEF
+    MATH_DIST_METHODDEF
     {"erf",             math_erf,       METH_O,         math_erf_doc},
     {"erfc",            math_erfc,      METH_O,         math_erfc_doc},
     {"exp",             math_exp,       METH_O,         math_exp_doc},