Repository: numpy/numpy
File: numpy/_core/src/multiarray/stringdtype/dtype.c (C)

Code before:
}

int
stringdtype_stableargsort_loop(
        PyArrayMethod_Context *context,
        char *const *data, const npy_intp *dimensions, const npy_intp *strides,
        NpyAuxData *transferdata)
{
    return stringdtype_wrap_argsort_loop(
            context, data, dimensions, strides, transferdata,
            &npy_amergesort_impl);
}

int
stringdtype_get_sort_loop(
        PyArrayMethod_Context *context,
        int aligned, int move_references,
        const npy_intp *strides,
        PyArrayMethod_StridedLoop **out_loop,
        NpyAuxData **out_transferdata,
        NPY_ARRAYMETHOD_FLAGS *flags)
{
    PyArrayMethod_SortParameters *parameters = (PyArrayMethod_SortParameters *)context->parameters;
    *flags |= NPY_METH_NO_FLOATINGPOINT_ERRORS;

    if (parameters->flags == NPY_SORT_STABLE) {
        *out_loop = (PyArrayMethod_StridedLoop *)stringdtype_stablesort_loop;
    }
    else if (parameters->flags == NPY_SORT_DEFAULT) {
        *out_loop = (PyArrayMethod_StridedLoop *)stringdtype_defaultsort_loop;
    }
    else {
        PyErr_SetString(PyExc_RuntimeError, "unsupported sort kind");
        return -1;
    }
    return 0;
}

int
stringdtype_get_argsort_loop(
        PyArrayMethod_Context *context,
        int aligned, int move_references,
        const npy_intp *strides,
        PyArrayMethod_StridedLoop **out_loop,
        NpyAuxData **out_transferdata,
        NPY_ARRAYMETHOD_FLAGS *flags)
{
    PyArrayMethod_SortParameters *parameters = (PyArrayMethod_SortParameters *)context->parameters;
    *flags |= NPY_METH_NO_FLOATINGPOINT_ERRORS;

    if (parameters->flags == NPY_SORT_STABLE) {

Review comment (refactor):
Since the two loops that we return are identical except for passing on different `NpyAuxData`, I think we can simplify this to have just one loop, but already set the auxdata here (a nice aspect of what `get_loop` can do!).

Code after:
        const npy_intp *strides,
        PyArrayMethod_StridedLoop **out_loop,
        NpyAuxData **out_transferdata,
        NPY_ARRAYMETHOD_FLAGS *flags)
{
    PyArrayMethod_SortParameters *parameters = (PyArrayMethod_SortParameters *)context->parameters;
    *flags |= NPY_METH_NO_FLOATINGPOINT_ERRORS;

    if (parameters->flags == NPY_SORT_STABLE
        || parameters->flags == NPY_SORT_DEFAULT) {
        *out_loop = (PyArrayMethod_StridedLoop *)stringdtype_wrap_argsort_loop;
    }
    else {
        PyErr_SetString(PyExc_RuntimeError, "unsupported sort kind");
        return -1;
    }
    return 0;
}

static PyObject *
stringdtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, PyObject *kwds)
{
    static char *kwargs_strs[] = {"coerce", "na_object", NULL};

    PyObject *na_object = NULL;
    int coerce = 1;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|$pO&:StringDType",
                                     kwargs_strs, &coerce,
                                     _not_NoValue, &na_object)) {
        return NULL;
    }

    return new_stringdtype_instance(na_object, coerce);
}

static void
stringdtype_dealloc(PyArray_StringDTypeObject *self)
{
    Py_XDECREF(self->na_object);
    // this can be null if an error happens while initializing an instance
    if (self->allocator != NULL) {
        NpyString_free_allocator(self->allocator);
    }
    PyMem_RawFree((char *)self->na_name.buf);
    PyMem_RawFree((char *)self->default_string.buf);
    PyArrayDescr_Type.tp_dealloc((PyObject *)self);
}

static PyObject *
stringdtype_repr(PyArray_StringDTypeObject *self)