Function Strs_subscript

python/stringzilla.c:2041–2176

* This returns a `Strs` object of a potentially different layout: * - `STRS_U32_TAPE_VIEW` input yields a `STRS_U32_TAPE_VIEW` for `step=1`, `STRS_FRAGMENTED` otherwise. * - `STRS_U64_TAPE_VIEW` input yields a `STRS_U64_TAPE_VIEW` for `step=1`, `STRS_FRAGMENTED` otherwise. * - `STRS_U32_TAPE` input yields a `STRS_U32_TAPE_VIEW` for `step=1`, `STRS_FRAGMENTED` otherwise. * - `STRS_U64_TAP

 *  - `STRS_FRAGMENTED` input yields a `STRS_FRAGMENTED` output.
 */
static PyObject *Strs_subscript(Strs *self, PyObject *key) {

    if (PyLong_Check(key)) { return Strs_getitem(self, PyLong_AsSsize_t(key)); }

    if (!PySlice_Check(key)) {
        PyErr_SetString(PyExc_TypeError, "Strs indices must be integers or slices");
        return NULL;
    }

    // Sanity checks
    Py_ssize_t count = Strs_len(self);
    Py_ssize_t start, stop, step;
    if (PySlice_Unpack(key, &start, &stop, &step) < 0) return NULL;
    Py_ssize_t result_count = PySlice_AdjustIndices(count, &start, &stop, step);
    if (result_count < 0) return NULL;

    // Create a new `Strs` object
    Strs *result = (Strs *)StrsType.tp_alloc(&StrsType, 0);
    if (result == NULL && PyErr_NoMemory()) return NULL;

    if (result_count == 0) {
        result->layout = STRS_FRAGMENTED;
        result->data.fragmented.count = 0;
        result->data.fragmented.spans = NULL;
        result->data.fragmented.parent = NULL;
        sz_memory_allocator_init_default(&result->data.fragmented.allocator);
        return (PyObject *)result;
    }

    // If a step is requested, we have to create a new `FRAGMENTED` instance of `Strs`,
    // even if the original one was a tape layout.
    if (step != 1) {
        sz_string_view_t *new_spans = (sz_string_view_t *)malloc(result_count * sizeof(sz_string_view_t));
        if (new_spans == NULL) {
            Py_XDECREF(result);
            PyErr_SetString(PyExc_MemoryError, "Unable to allocate memory for fragmented spans");
            return NULL;
        }

        get_string_at_offset_t getter = str_at_offset_getter(self);
        result->layout = STRS_FRAGMENTED;
        result->data.fragmented.count = result_count;
        result->data.fragmented.spans = new_spans;
        result->data.fragmented.parent = NULL;
        sz_memory_allocator_init_default(&result->data.fragmented.allocator);

        // Populate the new fragmented array using `get_string_at_offset`
        sz_size_t j = 0;
        if (step > 0)
            for (Py_ssize_t i = start; i < stop; i += step, ++j) {
                getter(self, i, count, &result->data.fragmented.parent, &new_spans[j].start, &new_spans[j].length);
            }
        else
            for (Py_ssize_t i = start; i > stop; i += step, ++j) {
                getter(self, i, count, &result->data.fragmented.parent, &new_spans[j].start, &new_spans[j].length);
            }

        // Ensure the parent string isn't prematurely deallocated by this view.