Add string comparison support so that x.sort() works

Author: peytondmurray

stringdtype/README.md

@@ -10,18 +10,26 @@ NumPy.
 Ensure Meson and NumPy are installed in the python environment you would like to use:
 
 ```
-$ python3 -m pip install meson meson-python numpy build patchelf
+$ python3 -m pip install meson meson-python build patchelf
 ```
 
-Build with meson, create a wheel, and install it
+It is important to have the latest development version of numpy installed.
+Nightly wheels work well for this purpose, and can be installed easily:
 
+```bash
+$ pip install -i https://pypi.anaconda.org/scipy-wheels-nightly/simple numpy
 ```
+
+Build with meson, create a wheel, and install it.
+
+```bash
 $ rm -r dist/
 $ meson build
 $ python -m build --wheel -Cbuilddir=build
-$ python -m pip install dist/asciidtype*.whl
 ```
 
-The `mesonpy` build backend for pip [does not currently support editable
-installs](https://github.com/mesonbuild/meson-python/issues/47), so `pip install
--e .` will not work.
+Or simply install directly, taking care to install without build isolation:
+
+```bash
+$ pip install -v . --no-build-isolation
+```

stringdtype/stringdtype/__init__.py

@@ -5,7 +5,6 @@
 from .scalar import StringScalar  # isort: skip
 from ._main import StringDType, _memory_usage
 
-
 __all__ = [
     "StringDType",
     "StringScalar",

stringdtype/stringdtype/src/dtype.c

@@ -155,6 +155,15 @@ stringdtype_getitem(StringDTypeObject *descr, char **dataptr)
     return res;
 }
 
+// Implementation of PyArray_CompareFunc.
+int
+compare_strings(char **a, char **b, PyArrayObject *arr)
+{
+    ss *ss_a = (ss *)*a;
+    ss *ss_b = (ss *)*b;
+    return strcmp(ss_a->buf, ss_b->buf);
+}
+
 static StringDTypeObject *
 stringdtype_ensure_canonical(StringDTypeObject *self)
 {
@@ -170,6 +179,7 @@ static PyType_Slot StringDType_Slots[] = {
         {NPY_DT_setitem, &stringdtype_setitem},
         {NPY_DT_getitem, &stringdtype_getitem},
         {NPY_DT_ensure_canonical, &stringdtype_ensure_canonical},
+        {NPY_DT_PyArray_ArrFuncs_compare, &compare_strings},
         {0, NULL}};
 
 static PyObject *
@@ -311,6 +321,7 @@ init_string_dtype(void)
     /* Loaded dynamically, so may need to be set here: */
     ((PyObject *)&StringDType)->ob_type = &PyArrayDTypeMeta_Type;
     ((PyTypeObject *)&StringDType)->tp_base = &PyArrayDescr_Type;
+
     if (PyType_Ready((PyTypeObject *)&StringDType) < 0) {
         return -1;
     }

stringdtype/stringdtype/src/main.c

@@ -90,7 +90,7 @@ PyInit__main(void)
     if (_import_array() < 0) {
         return NULL;
     }
-    if (import_experimental_dtype_api(7) < 0) {
+    if (import_experimental_dtype_api(8) < 0) {
         return NULL;
     }
 

stringdtype/tests/test_stringdtype.py

@@ -1,6 +1,7 @@
 import concurrent.futures
 import os
 import pickle
+import string
 import tempfile
 
 import numpy as np
@@ -14,6 +15,17 @@ def string_list():
     return ["abc", "def", "ghi"]
 
 
+@pytest.fixture
+def string_list_long():
+    abcs = string.ascii_lowercase
+
+    pairs = []
+    for pair in zip(abcs, abcs[1:] + abcs[0]):
+        pairs.append("".join(pair))
+
+    return pairs
+
+
 def test_scalar_creation():
     assert str(StringScalar("abc")) == "abc"
 
@@ -161,3 +173,13 @@ def test_pickle(string_list):
     assert res[1] == dtype
 
     os.remove(f.name)
+
+
+def test_sort(string_list_long):
+    """Test that sorting matches python's internal sorting."""
+    arr = np.array(string_list_long, dtype=StringDType())
+    arr_sorted = np.array(sorted(string_list_long), dtype=StringDType())
+
+    np.random.default_rng().shuffle(arr)
+    arr.sort()
+    np.testing.assert_array_equal(arr, arr_sorted)