Merge pull request #332 from jeromekelleher/prevent-state-updates

general_stat funcs get copies of underlying data.

Author: jeromekelleher

python/_tskitmodule.c

@@ -6503,10 +6503,11 @@ general_stat_func(size_t K, double *X, size_t M, double *Y, void *params)
     npy_intp X_dims = (npy_intp) K;
     npy_intp *Y_dims;
 
-    X_array = (PyArrayObject *) PyArray_SimpleNewFromData(1, &X_dims, NPY_FLOAT64, X);
+    X_array = (PyArrayObject *) PyArray_SimpleNew(1, &X_dims, NPY_FLOAT64);
     if (X_array == NULL) {
         goto out;
     }
+    memcpy(PyArray_DATA(X_array), X, K * sizeof(*X));
     arglist = Py_BuildValue("(O)", X_array);
     if (arglist == NULL) {
         goto out;

python/tests/test_tree_stats.py

@@ -3315,6 +3315,25 @@ def get_tree_sequence(self):
                               mutation_rate=2, random_seed=1)
         return ts
 
+    def test_function_cannot_update_state(self):
+        ts = self.get_tree_sequence()
+
+        def f(x):
+            out = x.copy()
+            x[:] = 0.0
+            return out
+
+        def g(x):
+            return x
+
+        x = ts.sample_count_stat(
+            [ts.samples()], f, output_dim=1, strict=False, mode="node",
+            span_normalise=False)
+        y = ts.sample_count_stat(
+            [ts.samples()], g, output_dim=1, strict=False, mode="node",
+            span_normalise=False)
+        self.assertArrayEqual(x, y)
+
     def test_default_mode(self):
         ts = msprime.simulate(10, recombination_rate=1, random_seed=2)
         W = np.ones((ts.num_samples, 2))