diff options
Diffstat (limited to 'Objects/stringlib/join.h')
| -rw-r--r-- | Objects/stringlib/join.h | 132 | 
1 files changed, 132 insertions, 0 deletions
| diff --git a/Objects/stringlib/join.h b/Objects/stringlib/join.h new file mode 100644 index 0000000000..d1d6e532c5 --- /dev/null +++ b/Objects/stringlib/join.h @@ -0,0 +1,132 @@ +/* stringlib: bytes joining implementation */ + +#if STRINGLIB_SIZEOF_CHAR != 1 +#error join.h only compatible with byte-wise strings +#endif + +Py_LOCAL_INLINE(PyObject *) +STRINGLIB(bytes_join)(PyObject *sep, PyObject *iterable) +{ +    char *sepstr = STRINGLIB_STR(sep); +    const Py_ssize_t seplen = STRINGLIB_LEN(sep); +    PyObject *res = NULL; +    char *p; +    Py_ssize_t seqlen = 0; +    Py_ssize_t sz = 0; +    Py_ssize_t i, nbufs; +    PyObject *seq, *item; +    Py_buffer *buffers = NULL; +#define NB_STATIC_BUFFERS 10 +    Py_buffer static_buffers[NB_STATIC_BUFFERS]; + +    seq = PySequence_Fast(iterable, "can only join an iterable"); +    if (seq == NULL) { +        return NULL; +    } + +    seqlen = PySequence_Fast_GET_SIZE(seq); +    if (seqlen == 0) { +        Py_DECREF(seq); +        return STRINGLIB_NEW(NULL, 0); +    } +#ifndef STRINGLIB_MUTABLE +    if (seqlen == 1) { +        item = PySequence_Fast_GET_ITEM(seq, 0); +        if (STRINGLIB_CHECK_EXACT(item)) { +            Py_INCREF(item); +            Py_DECREF(seq); +            return item; +        } +    } +#endif +    if (seqlen > NB_STATIC_BUFFERS) { +        buffers = PyMem_NEW(Py_buffer, seqlen); +        if (buffers == NULL) { +            Py_DECREF(seq); +            return NULL; +        } +    } +    else { +        buffers = static_buffers; +    } + +    /* Here is the general case.  Do a pre-pass to figure out the total +     * amount of space we'll need (sz), and see whether all arguments are +     * buffer-compatible. +     */ +    for (i = 0, nbufs = 0; i < seqlen; i++) { +        Py_ssize_t itemlen; +        item = PySequence_Fast_GET_ITEM(seq, i); +        if (_getbuffer(item, &buffers[i]) < 0) { +            PyErr_Format(PyExc_TypeError, +                         "sequence item %zd: expected bytes, bytearray, " +                         "or an object with the buffer interface, %.80s found", +                         i, Py_TYPE(item)->tp_name); +            goto error; +        } +        nbufs = i + 1;  /* for error cleanup */ +        itemlen = buffers[i].len; +        if (itemlen > PY_SSIZE_T_MAX - sz) { +            PyErr_SetString(PyExc_OverflowError, +                            "join() result is too long"); +            goto error; +        } +        sz += itemlen; +        if (i != 0) { +            if (seplen > PY_SSIZE_T_MAX - sz) { +                PyErr_SetString(PyExc_OverflowError, +                                "join() result is too long"); +                goto error; +            } +            sz += seplen; +        } +        if (seqlen != PySequence_Fast_GET_SIZE(seq)) { +            PyErr_SetString(PyExc_RuntimeError, +                            "sequence changed size during iteration"); +            goto error; +        } +    } + +    /* Allocate result space. */ +    res = STRINGLIB_NEW(NULL, sz); +    if (res == NULL) +        goto error; + +    /* Catenate everything. */ +    p = STRINGLIB_STR(res); +    if (!seplen) { +        /* fast path */ +        for (i = 0; i < nbufs; i++) { +            Py_ssize_t n = buffers[i].len; +            char *q = buffers[i].buf; +            Py_MEMCPY(p, q, n); +            p += n; +        } +        goto done; +    } +    for (i = 0; i < nbufs; i++) { +        Py_ssize_t n; +        char *q; +        if (i) { +            Py_MEMCPY(p, sepstr, seplen); +            p += seplen; +        } +        n = buffers[i].len; +        q = buffers[i].buf; +        Py_MEMCPY(p, q, n); +        p += n; +    } +    goto done; + +error: +    res = NULL; +done: +    Py_DECREF(seq); +    for (i = 0; i < nbufs; i++) +        PyBuffer_Release(&buffers[i]); +    if (buffers != static_buffers) +        PyMem_FREE(buffers); +    return res; +} + +#undef NB_STATIC_BUFFERS | 
