Merge pull request #1892 from jeromekelleher/better-error-general-stat

Better error message in general stat

Author: mergify[bot]

python/_tskitmodule.c

@@ -8192,13 +8192,23 @@ general_stat_func(tsk_size_t K, const double *X, tsk_size_t M, double *Y, void *
         goto out;
     }
     Y_array = (PyArrayObject *) PyArray_FromAny(
-        result, PyArray_DescrFromType(NPY_FLOAT64), 1, 1, NPY_ARRAY_IN_ARRAY, NULL);
+        result, PyArray_DescrFromType(NPY_FLOAT64), 0, 0, NPY_ARRAY_IN_ARRAY, NULL);
     if (Y_array == NULL) {
         goto out;
     }
+    if (PyArray_NDIM(Y_array) != 1) {
+        PyErr_Format(PyExc_ValueError,
+            "Array returned by general_stat callback is %d dimensional; "
+            "must be 1D",
+            (int) PyArray_NDIM(Y_array));
+        goto out;
+    }
     Y_dims = PyArray_DIMS(Y_array);
     if (Y_dims[0] != (npy_intp) M) {
-        PyErr_SetString(PyExc_ValueError, "Incorrect callback output dimensions");
+        PyErr_Format(PyExc_ValueError,
+            "Array returned by general_stat callback is of length %d; "
+            "must be %d",
+            Y_dims[0], M);
         goto out;
     }
     /* Copy the contents of the return Y array into Y */

python/tests/test_tree_stats.py

@@ -6150,3 +6150,38 @@ def test_uncalibrated_time_general_stat(self, ts_fixture):
                 W.shape[1],
                 mode="branch",
             )
+
+
+class TestGeneralStatCallbackErrors:
+    def test_zero_d(self, ts_fixture):
+        def f_0d(_):
+            return 0
+
+        msg = "Array returned by general_stat callback is 0 dimensional; must be 1D"
+        with pytest.raises(ValueError, match=msg):
+            ts_fixture.sample_count_stat(
+                sample_sets=[ts_fixture.samples()], f=f_0d, output_dim=1, strict=False
+            )
+
+    def test_two_d(self, ts_fixture):
+        def f_2d(x):
+            return np.array([x])
+
+        msg = "Array returned by general_stat callback is 2 dimensional; must be 1D"
+        with pytest.raises(ValueError, match=msg):
+            ts_fixture.sample_count_stat(
+                sample_sets=[ts_fixture.samples()], f=f_2d, output_dim=1, strict=False
+            )
+
+    def test_wrong_length(self, ts_fixture):
+        def f_too_long(_):
+            return np.array([0, 0])
+
+        msg = "Array returned by general_stat callback is of length 2; must be 1"
+        with pytest.raises(ValueError, match=msg):
+            ts_fixture.sample_count_stat(
+                sample_sets=[ts_fixture.samples()],
+                f=f_too_long,
+                output_dim=1,
+                strict=False,
+            )