Switch to new numba string type

Author: benjeffery

python/_tskitmodule.c

@@ -24,7 +24,8 @@
  */
 
 #define PY_SSIZE_T_CLEAN
-#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#define NPY_NO_DEPRECATED_API NPY_2_0_API_VERSION
+#define NPY_TARGET_VERSION NPY_2_0_API_VERSION
 #define TSK_BUG_ASSERT_MESSAGE                                                          \
     "Please open an issue on"                                                           \
     " GitHub, ideally with a reproducible example."                                     \
@@ -10779,110 +10780,56 @@ TreeSequence_decode_ragged_string_column(
     TreeSequence *self, tsk_size_t num_rows, const char *data, const tsk_size_t *offset)
 {
     PyObject *ret = NULL;
-    PyArrayObject *ret_array = NULL;
+    PyObject *array = NULL;
+    PyArray_StringDTypeObject *string_dtype = NULL;
+    npy_string_allocator *allocator = NULL;
+    char *array_data = NULL;
     npy_intp dims[1];
-    char *string_data = NULL;
     tsk_size_t i;
-    tsk_size_t max_length;
-    tsk_size_t len;
-    tsk_size_t start;
-    tsk_size_t end;
-    char *dest;
-    bool all_length_one;
+    int pack_result;
 
+    string_dtype = (PyArray_StringDTypeObject *) PyArray_DescrFromType(NPY_VSTRING);
+    if (string_dtype == NULL) {
+        goto out;
+    }
     dims[0] = (npy_intp) num_rows;
-
-    /* If all strings are length one we can back the array with tskit memory */
-    all_length_one = true;
-    for (i = 0; i < num_rows; i++) {
-        len = offset[i + 1] - offset[i];
-        if (len != 1) {
-            all_length_one = false;
-            break;
-        }
+    array = PyArray_Zeros(1, dims, (PyArray_Descr *) string_dtype, 0);
+    if (array == NULL) {
+        goto out;
+    }
+    allocator = NpyString_acquire_allocator(string_dtype);
+    if (allocator == NULL) {
+        goto out;
     }
+    array_data = (char *) PyArray_DATA((PyArrayObject *) array);
+    for (i = 0; i < num_rows; i++) {
 
-    if (all_length_one && num_rows > 0) {
-        ret_array = (PyArrayObject *) PyArray_New(&PyArray_Type, /* subtype */
-            1,                                                   /* nd */
-            dims,                                                /* dims */
-            NPY_STRING,                                          /* type_num */
-            NULL,                                                /* strides */
-            (void *) data,                                       /* data */
-            1,   /* itemsize = 1 for S1 */
-            0,   /* flags */
-            NULL /* obj */
-        );
-        if (ret_array == NULL) {
-            goto out;
-        }
+        pack_result = NpyString_pack(allocator,
+            (npy_packed_static_string
+                    *) (array_data + (i * ((PyArray_Descr *) string_dtype)->elsize)),
+            data + offset[i], offset[i + 1] - offset[i]);
 
-        if (PyArray_SetBaseObject(ret_array, (PyObject *) self) != 0) {
+        if (pack_result == -1) {
             goto out;
         }
-        Py_INCREF(self);
-
-    } else {
-        /* We have to pad the strings to make a rectangular array */
-        max_length = 0;
-        for (i = 0; i < num_rows; i++) {
-            len = offset[i + 1] - offset[i];
-            if (len > max_length) {
-                max_length = len;
-            }
-        }
-
-        /* Ensure at least S1 */
-        if (max_length == 0) {
-            max_length = 1;
-        }
-
-        if (num_rows > 0) {
-            tsk_size_t total_size = num_rows * max_length;
-            string_data = PyMem_Calloc(total_size, sizeof(char));
-            if (string_data == NULL) {
-                PyErr_NoMemory();
-                goto out;
-            }
-
-            for (i = 0; i < num_rows; i++) {
-                start = offset[i];
-                end = offset[i + 1];
-                len = end - start;
-
-                dest = string_data + (i * max_length);
-                if (len > 0) {
-                    memcpy(dest, data + start, len);
-                }
-                /* PyMem_Calloc zero-fills so we don't need to pad */
-            }
-        }
+    }
 
-        ret_array = (PyArrayObject *) PyArray_New(&PyArray_Type, /* subtype */
-            1,                                                   /* nd */
-            dims,                                                /* dims */
-            NPY_STRING,                                          /* type_num */
-            NULL,                                                /* strides */
-            string_data,                                         /* data */
-            (int) max_length,                                    /* itemsize */
-            NPY_ARRAY_OWNDATA, /* flags - we own the data */
-            NULL               /* obj */
-        );
-        if (ret_array == NULL) {
-            goto out;
-        }
-        string_data = NULL; /* Array now owns the memory */
+    if (PyArray_SetBaseObject((PyArrayObject *) array, (PyObject *) self) != 0) {
+        goto out;
     }
+    Py_INCREF(self);
 
-    /* Make array read-only for consistent semantics with tskit memory case */
-    PyArray_CLEARFLAGS(ret_array, NPY_ARRAY_WRITEABLE);
+    /* Clear the writeable flag to match other arrays semantics */
+    PyArray_CLEARFLAGS((PyArrayObject *) array, NPY_ARRAY_WRITEABLE);
 
-    ret = (PyObject *) ret_array;
-    ret_array = NULL;
+    ret = array;
+    array = NULL;
 
 out:
-    PyMem_Free(string_data);
-    Py_XDECREF(ret_array);
+    if (allocator != NULL) {
+        NpyString_release_allocator(allocator);
+    }
+    Py_XDECREF(array);
     return ret;
 }
 
@@ -14705,11 +14652,16 @@ static struct PyModuleDef tskitmodule = {
 PyObject *
 PyInit__tskit(void)
 {
-    PyObject *module = PyModule_Create(&tskitmodule);
+    PyObject *module;
+
+    if (PyArray_ImportNumPyAPI() < 0) {
+        return NULL;
+    }
+
+    module = PyModule_Create(&tskitmodule);
     if (module == NULL) {
         return NULL;
     }
-    import_array();
 
     if (register_lwt_class(module) != 0) {
         return NULL;

python/pyproject.toml

@@ -43,7 +43,7 @@ keywords = [
 requires-python = ">=3.9"
 dependencies = [
     "jsonschema>=3.0.0",
-    "numpy>=1.23.5",
+    "numpy>2",
 ]
 
 [project.urls]

python/requirements/development.txt

@@ -17,7 +17,7 @@ msprime>=1.0.0
 networkx
 newick
 ninja
-numpy
+numpy>=2
 packaging
 portion
 pre-commit

python/requirements/development.yml

@@ -22,7 +22,7 @@ dependencies:
   - msprime>=1.0.0
   - networkx
   - ninja
-  - numpy<2
+  - numpy>2
   - packaging
   - portion
   - pre-commit

python/tests/test_lowlevel.py

@@ -1916,8 +1916,9 @@ def test_array_read_only(self, name, ts_fixture):
             a[:] = 0
         with pytest.raises(ValueError, match="assignment destination"):
             a[0] = 0
-        with pytest.raises(ValueError, match="cannot set WRITEABLE"):
-            a.setflags(write=True)
+        if name != "sites_ancestral_state":
+            with pytest.raises(ValueError, match="cannot set WRITEABLE"):
+                a.setflags(write=True)
 
     @pytest.mark.parametrize("name", ARRAY_NAMES)
     def test_array_properties(self, name, ts_fixture):
@@ -1927,13 +1928,17 @@ def test_array_properties(self, name, ts_fixture):
         assert not a.flags.writeable
         assert a.flags.aligned
         assert a.flags.c_contiguous
-        assert not a.flags.owndata
+        if name == "sites_ancestral_state":
+            assert a.flags.owndata
+        else:
+            assert not a.flags.owndata
         b = getattr(ts_fixture, name)
         assert a is not b
         assert np.all(a == b)
         # This checks that the underlying pointer to memory is the same in
         # both arrays.
-        assert a.__array_interface__ == b.__array_interface__
+        if name != "sites_ancestral_state":
+            assert a.__array_interface__ == b.__array_interface__
 
     @pytest.mark.parametrize("name", ARRAY_NAMES)
     def test_array_lifetime(self, name, ts_fixture):
@@ -1976,7 +1981,8 @@ def test_individuals_nodes(self, ts_fixture):
         assert a3 is not a2
 
     @pytest.mark.parametrize(
-        "site_lengths", ["none", "all-0", "all-1", "all-2", "mixed", "very_long"]
+        "site_lengths",
+        ["none", "all-0", "all-1", "all-2", "mixed", "very_long", "unicode"],
     )
     def test_sites_ancestral_state(self, ts_fixture, site_lengths):
         if site_lengths == "none":
@@ -1995,6 +2001,7 @@ def test_sites_ancestral_state(self, ts_fixture, site_lengths):
                     "all-2": lambda i, site: "YO",
                     "mixed": lambda i, site: "A" * i,
                     "very_long": lambda i, site: "A" * 100_000_000 if i == 0 else "A",
+                    "unicode": lambda i, site: "💩" * (i + 1),
                 }
 
                 get_ancestral_state = ancestral_state_map[site_lengths]
@@ -2003,15 +2010,22 @@ def test_sites_ancestral_state(self, ts_fixture, site_lengths):
                         site.replace(ancestral_state=get_ancestral_state(i, site))
                     )
                 ts = tables.tree_sequence()
-        ts = ts.ll_tree_sequence
+        ll_ts = ts.ll_tree_sequence
+
+        a = ll_ts.sites_ancestral_state
+        # Contents
+        if site_lengths == "none":
+            assert a.size == 0
+        else:
+            for site in ts.sites():
+                assert a[site.id] == site.ancestral_state
 
         # Read only
         with pytest.raises(AttributeError, match="not writable"):
-            ts.sites_ancestral_state = None
+            ll_ts.sites_ancestral_state = None
         with pytest.raises(AttributeError, match="not writable"):
-            del ts.sites_ancestral_state
+            del ll_ts.sites_ancestral_state
 
-        a = ts.sites_ancestral_state
         with pytest.raises(ValueError, match="assignment destination"):
             a[:] = 0
         with pytest.raises(ValueError, match="assignment destination"):
@@ -2021,18 +2035,17 @@ def test_sites_ancestral_state(self, ts_fixture, site_lengths):
                 a.setflags(write=True)
 
         # Properties
-        a = ts.sites_ancestral_state
         assert a.flags.aligned
         assert a.flags.c_contiguous
-        b = ts.sites_ancestral_state
+        b = ll_ts.sites_ancestral_state
         assert a is not b
         assert np.all(a == b)
 
         # Lifetime
-        a1 = ts.sites_ancestral_state
+        a1 = ll_ts.sites_ancestral_state
         a2 = a1.copy()
         assert a1 is not a2
-        del ts
+        del ll_ts
         # Do some memory operations
         a3 = np.ones(10**6)
         assert np.all(a1 == a2)