bytearrayobject.c | searchcode

/Objects/bytearrayobject.c

Large files files are truncated, but you can click here to view the full file
   1/* PyBytes (bytearray) implementation */
   2
   3#define PY_SSIZE_T_CLEAN
   4#include "Python.h"
   5#include "structmember.h"
   6#include "bytes_methods.h"
   7
   8static PyByteArrayObject *nullbytes = NULL;
   9
  10void
  11PyByteArray_Fini(void)
  12{
  13    Py_CLEAR(nullbytes);
  14}
  15
  16int
  17PyByteArray_Init(void)
  18{
  19    nullbytes = PyObject_New(PyByteArrayObject, &PyByteArray_Type);
  20    if (nullbytes == NULL)
  21        return 0;
  22    nullbytes->ob_bytes = NULL;
  23    Py_SIZE(nullbytes) = nullbytes->ob_alloc = 0;
  24    nullbytes->ob_exports = 0;
  25    return 1;
  26}
  27
  28/* end nullbytes support */
  29
  30/* Helpers */
  31
  32static int
  33_getbytevalue(PyObject* arg, int *value)
  34{
  35    long face_value;
  36
  37    if (PyBytes_CheckExact(arg)) {
  38        if (Py_SIZE(arg) != 1) {
  39            PyErr_SetString(PyExc_ValueError, "string must be of size 1");
  40            return 0;
  41        }
  42        *value = Py_CHARMASK(((PyBytesObject*)arg)->ob_sval[0]);
  43        return 1;
  44    }
  45    else if (PyInt_Check(arg) || PyLong_Check(arg)) {
  46        face_value = PyLong_AsLong(arg);
  47    }
  48    else {
  49        PyObject *index = PyNumber_Index(arg);
  50        if (index == NULL) {
  51            PyErr_Format(PyExc_TypeError,
  52                         "an integer or string of size 1 is required");
  53            return 0;
  54        }
  55        face_value = PyLong_AsLong(index);
  56        Py_DECREF(index);
  57    }
  58
  59    if (face_value < 0 || face_value >= 256) {
  60        /* this includes the OverflowError in case the long is too large */
  61        PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
  62        return 0;
  63    }
  64
  65    *value = face_value;
  66    return 1;
  67}
  68
  69static Py_ssize_t
  70bytes_buffer_getreadbuf(PyByteArrayObject *self, Py_ssize_t index, const void **ptr)
  71{
  72    if ( index != 0 ) {
  73        PyErr_SetString(PyExc_SystemError,
  74                "accessing non-existent bytes segment");
  75        return -1;
  76    }
  77    *ptr = (void *)self->ob_bytes;
  78    return Py_SIZE(self);
  79}
  80
  81static Py_ssize_t
  82bytes_buffer_getwritebuf(PyByteArrayObject *self, Py_ssize_t index, const void **ptr)
  83{
  84    if ( index != 0 ) {
  85        PyErr_SetString(PyExc_SystemError,
  86                "accessing non-existent bytes segment");
  87        return -1;
  88    }
  89    *ptr = (void *)self->ob_bytes;
  90    return Py_SIZE(self);
  91}
  92
  93static Py_ssize_t
  94bytes_buffer_getsegcount(PyByteArrayObject *self, Py_ssize_t *lenp)
  95{
  96    if ( lenp )
  97        *lenp = Py_SIZE(self);
  98    return 1;
  99}
 100
 101static Py_ssize_t
 102bytes_buffer_getcharbuf(PyByteArrayObject *self, Py_ssize_t index, const char **ptr)
 103{
 104    if ( index != 0 ) {
 105        PyErr_SetString(PyExc_SystemError,
 106                "accessing non-existent bytes segment");
 107        return -1;
 108    }
 109    *ptr = self->ob_bytes;
 110    return Py_SIZE(self);
 111}
 112
 113static int
 114bytes_getbuffer(PyByteArrayObject *obj, Py_buffer *view, int flags)
 115{
 116        int ret;
 117        void *ptr;
 118        if (view == NULL) {
 119                obj->ob_exports++;
 120                return 0;
 121        }
 122        if (obj->ob_bytes == NULL)
 123                ptr = "";
 124        else
 125                ptr = obj->ob_bytes;
 126        ret = PyBuffer_FillInfo(view, (PyObject*)obj, ptr, Py_SIZE(obj), 0, flags);
 127        if (ret >= 0) {
 128                obj->ob_exports++;
 129        }
 130        return ret;
 131}
 132
 133static void
 134bytes_releasebuffer(PyByteArrayObject *obj, Py_buffer *view)
 135{
 136        obj->ob_exports--;
 137}
 138
 139static Py_ssize_t
 140_getbuffer(PyObject *obj, Py_buffer *view)
 141{
 142    PyBufferProcs *buffer = Py_TYPE(obj)->tp_as_buffer;
 143
 144    if (buffer == NULL || buffer->bf_getbuffer == NULL)
 145    {
 146        PyErr_Format(PyExc_TypeError,
 147                     "Type %.100s doesn't support the buffer API",
 148                     Py_TYPE(obj)->tp_name);
 149        return -1;
 150    }
 151
 152    if (buffer->bf_getbuffer(obj, view, PyBUF_SIMPLE) < 0)
 153            return -1;
 154    return view->len;
 155}
 156
 157static int
 158_canresize(PyByteArrayObject *self)
 159{
 160    if (self->ob_exports > 0) {
 161        PyErr_SetString(PyExc_BufferError,
 162                "Existing exports of data: object cannot be re-sized");
 163        return 0;
 164    }
 165    return 1;
 166}
 167
 168/* Direct API functions */
 169
 170PyObject *
 171PyByteArray_FromObject(PyObject *input)
 172{
 173    return PyObject_CallFunctionObjArgs((PyObject *)&PyByteArray_Type,
 174                                        input, NULL);
 175}
 176
 177PyObject *
 178PyByteArray_FromStringAndSize(const char *bytes, Py_ssize_t size)
 179{
 180    PyByteArrayObject *new;
 181    Py_ssize_t alloc;
 182
 183    if (size < 0) {
 184        PyErr_SetString(PyExc_SystemError,
 185            "Negative size passed to PyByteArray_FromStringAndSize");
 186        return NULL;
 187    }
 188
 189    new = PyObject_New(PyByteArrayObject, &PyByteArray_Type);
 190    if (new == NULL)
 191        return NULL;
 192
 193    if (size == 0) {
 194        new->ob_bytes = NULL;
 195        alloc = 0;
 196    }
 197    else {
 198        alloc = size + 1;
 199        new->ob_bytes = PyMem_Malloc(alloc);
 200        if (new->ob_bytes == NULL) {
 201            Py_DECREF(new);
 202            return PyErr_NoMemory();
 203        }
 204        if (bytes != NULL)
 205            memcpy(new->ob_bytes, bytes, size);
 206        new->ob_bytes[size] = '\0';  /* Trailing null byte */
 207    }
 208    Py_SIZE(new) = size;
 209    new->ob_alloc = alloc;
 210    new->ob_exports = 0;
 211
 212    return (PyObject *)new;
 213}
 214
 215Py_ssize_t
 216PyByteArray_Size(PyObject *self)
 217{
 218    assert(self != NULL);
 219    assert(PyByteArray_Check(self));
 220
 221    return PyByteArray_GET_SIZE(self);
 222}
 223
 224char  *
 225PyByteArray_AsString(PyObject *self)
 226{
 227    assert(self != NULL);
 228    assert(PyByteArray_Check(self));
 229
 230    return PyByteArray_AS_STRING(self);
 231}
 232
 233int
 234PyByteArray_Resize(PyObject *self, Py_ssize_t size)
 235{
 236    void *sval;
 237    Py_ssize_t alloc = ((PyByteArrayObject *)self)->ob_alloc;
 238
 239    assert(self != NULL);
 240    assert(PyByteArray_Check(self));
 241    assert(size >= 0);
 242
 243    if (size == Py_SIZE(self)) {
 244        return 0;
 245    }
 246    if (!_canresize((PyByteArrayObject *)self)) {
 247        return -1;
 248    }
 249
 250    if (size < alloc / 2) {
 251        /* Major downsize; resize down to exact size */
 252        alloc = size + 1;
 253    }
 254    else if (size < alloc) {
 255        /* Within allocated size; quick exit */
 256        Py_SIZE(self) = size;
 257        ((PyByteArrayObject *)self)->ob_bytes[size] = '\0'; /* Trailing null */
 258        return 0;
 259    }
 260    else if (size <= alloc * 1.125) {
 261        /* Moderate upsize; overallocate similar to list_resize() */
 262        alloc = size + (size >> 3) + (size < 9 ? 3 : 6);
 263    }
 264    else {
 265        /* Major upsize; resize up to exact size */
 266        alloc = size + 1;
 267    }
 268
 269    sval = PyMem_Realloc(((PyByteArrayObject *)self)->ob_bytes, alloc);
 270    if (sval == NULL) {
 271        PyErr_NoMemory();
 272        return -1;
 273    }
 274
 275    ((PyByteArrayObject *)self)->ob_bytes = sval;
 276    Py_SIZE(self) = size;
 277    ((PyByteArrayObject *)self)->ob_alloc = alloc;
 278    ((PyByteArrayObject *)self)->ob_bytes[size] = '\0'; /* Trailing null byte */
 279
 280    return 0;
 281}
 282
 283PyObject *
 284PyByteArray_Concat(PyObject *a, PyObject *b)
 285{
 286    Py_ssize_t size;
 287    Py_buffer va, vb;
 288    PyByteArrayObject *result = NULL;
 289
 290    va.len = -1;
 291    vb.len = -1;
 292    if (_getbuffer(a, &va) < 0  ||
 293        _getbuffer(b, &vb) < 0) {
 294            PyErr_Format(PyExc_TypeError, "can't concat %.100s to %.100s",
 295                         Py_TYPE(a)->tp_name, Py_TYPE(b)->tp_name);
 296            goto done;
 297    }
 298
 299    size = va.len + vb.len;
 300    if (size < 0) {
 301            return PyErr_NoMemory();
 302            goto done;
 303    }
 304
 305    result = (PyByteArrayObject *) PyByteArray_FromStringAndSize(NULL, size);
 306    if (result != NULL) {
 307        memcpy(result->ob_bytes, va.buf, va.len);
 308        memcpy(result->ob_bytes + va.len, vb.buf, vb.len);
 309    }
 310
 311  done:
 312    if (va.len != -1)
 313        PyBuffer_Release(&va);
 314    if (vb.len != -1)
 315        PyBuffer_Release(&vb);
 316    return (PyObject *)result;
 317}
 318
 319/* Functions stuffed into the type object */
 320
 321static Py_ssize_t
 322bytes_length(PyByteArrayObject *self)
 323{
 324    return Py_SIZE(self);
 325}
 326
 327static PyObject *
 328bytes_iconcat(PyByteArrayObject *self, PyObject *other)
 329{
 330    Py_ssize_t mysize;
 331    Py_ssize_t size;
 332    Py_buffer vo;
 333
 334    if (_getbuffer(other, &vo) < 0) {
 335        PyErr_Format(PyExc_TypeError, "can't concat %.100s to %.100s",
 336                     Py_TYPE(other)->tp_name, Py_TYPE(self)->tp_name);
 337        return NULL;
 338    }
 339
 340    mysize = Py_SIZE(self);
 341    size = mysize + vo.len;
 342    if (size < 0) {
 343        PyBuffer_Release(&vo);
 344        return PyErr_NoMemory();
 345    }
 346    if (size < self->ob_alloc) {
 347        Py_SIZE(self) = size;
 348        self->ob_bytes[Py_SIZE(self)] = '\0'; /* Trailing null byte */
 349    }
 350    else if (PyByteArray_Resize((PyObject *)self, size) < 0) {
 351        PyBuffer_Release(&vo);
 352        return NULL;
 353    }
 354    memcpy(self->ob_bytes + mysize, vo.buf, vo.len);
 355    PyBuffer_Release(&vo);
 356    Py_INCREF(self);
 357    return (PyObject *)self;
 358}
 359
 360static PyObject *
 361bytes_repeat(PyByteArrayObject *self, Py_ssize_t count)
 362{
 363    PyByteArrayObject *result;
 364    Py_ssize_t mysize;
 365    Py_ssize_t size;
 366
 367    if (count < 0)
 368        count = 0;
 369    mysize = Py_SIZE(self);
 370    size = mysize * count;
 371    if (count != 0 && size / count != mysize)
 372        return PyErr_NoMemory();
 373    result = (PyByteArrayObject *)PyByteArray_FromStringAndSize(NULL, size);
 374    if (result != NULL && size != 0) {
 375        if (mysize == 1)
 376            memset(result->ob_bytes, self->ob_bytes[0], size);
 377        else {
 378            Py_ssize_t i;
 379            for (i = 0; i < count; i++)
 380                memcpy(result->ob_bytes + i*mysize, self->ob_bytes, mysize);
 381        }
 382    }
 383    return (PyObject *)result;
 384}
 385
 386static PyObject *
 387bytes_irepeat(PyByteArrayObject *self, Py_ssize_t count)
 388{
 389    Py_ssize_t mysize;
 390    Py_ssize_t size;
 391
 392    if (count < 0)
 393        count = 0;
 394    mysize = Py_SIZE(self);
 395    size = mysize * count;
 396    if (count != 0 && size / count != mysize)
 397        return PyErr_NoMemory();
 398    if (size < self->ob_alloc) {
 399        Py_SIZE(self) = size;
 400        self->ob_bytes[Py_SIZE(self)] = '\0'; /* Trailing null byte */
 401    }
 402    else if (PyByteArray_Resize((PyObject *)self, size) < 0)
 403        return NULL;
 404
 405    if (mysize == 1)
 406        memset(self->ob_bytes, self->ob_bytes[0], size);
 407    else {
 408        Py_ssize_t i;
 409        for (i = 1; i < count; i++)
 410            memcpy(self->ob_bytes + i*mysize, self->ob_bytes, mysize);
 411    }
 412
 413    Py_INCREF(self);
 414    return (PyObject *)self;
 415}
 416
 417static PyObject *
 418bytes_getitem(PyByteArrayObject *self, Py_ssize_t i)
 419{
 420    if (i < 0)
 421        i += Py_SIZE(self);
 422    if (i < 0 || i >= Py_SIZE(self)) {
 423        PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
 424        return NULL;
 425    }
 426    return PyInt_FromLong((unsigned char)(self->ob_bytes[i]));
 427}
 428
 429static PyObject *
 430bytes_subscript(PyByteArrayObject *self, PyObject *index)
 431{
 432    if (PyIndex_Check(index)) {
 433        Py_ssize_t i = PyNumber_AsSsize_t(index, PyExc_IndexError);
 434
 435        if (i == -1 && PyErr_Occurred())
 436            return NULL;
 437
 438        if (i < 0)
 439            i += PyByteArray_GET_SIZE(self);
 440
 441        if (i < 0 || i >= Py_SIZE(self)) {
 442            PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
 443            return NULL;
 444        }
 445        return PyInt_FromLong((unsigned char)(self->ob_bytes[i]));
 446    }
 447    else if (PySlice_Check(index)) {
 448        Py_ssize_t start, stop, step, slicelength, cur, i;
 449        if (PySlice_GetIndicesEx((PySliceObject *)index,
 450                                 PyByteArray_GET_SIZE(self),
 451                                 &start, &stop, &step, &slicelength) < 0) {
 452            return NULL;
 453        }
 454
 455        if (slicelength <= 0)
 456            return PyByteArray_FromStringAndSize("", 0);
 457        else if (step == 1) {
 458            return PyByteArray_FromStringAndSize(self->ob_bytes + start,
 459                                             slicelength);
 460        }
 461        else {
 462            char *source_buf = PyByteArray_AS_STRING(self);
 463            char *result_buf = (char *)PyMem_Malloc(slicelength);
 464            PyObject *result;
 465
 466            if (result_buf == NULL)
 467                return PyErr_NoMemory();
 468
 469            for (cur = start, i = 0; i < slicelength;
 470                 cur += step, i++) {
 471                     result_buf[i] = source_buf[cur];
 472            }
 473            result = PyByteArray_FromStringAndSize(result_buf, slicelength);
 474            PyMem_Free(result_buf);
 475            return result;
 476        }
 477    }
 478    else {
 479        PyErr_SetString(PyExc_TypeError, "bytearray indices must be integers");
 480        return NULL;
 481    }
 482}
 483
 484static int
 485bytes_setslice(PyByteArrayObject *self, Py_ssize_t lo, Py_ssize_t hi,
 486               PyObject *values)
 487{
 488    Py_ssize_t avail, needed;
 489    void *bytes;
 490    Py_buffer vbytes;
 491    int res = 0;
 492
 493    vbytes.len = -1;
 494    if (values == (PyObject *)self) {
 495        /* Make a copy and call this function recursively */
 496        int err;
 497        values = PyByteArray_FromObject(values);
 498        if (values == NULL)
 499            return -1;
 500        err = bytes_setslice(self, lo, hi, values);
 501        Py_DECREF(values);
 502        return err;
 503    }
 504    if (values == NULL) {
 505        /* del b[lo:hi] */
 506        bytes = NULL;
 507        needed = 0;
 508    }
 509    else {
 510            if (_getbuffer(values, &vbytes) < 0) {
 511                    PyErr_Format(PyExc_TypeError,
 512                                 "can't set bytearray slice from %.100s",
 513                                 Py_TYPE(values)->tp_name);
 514                    return -1;
 515            }
 516            needed = vbytes.len;
 517            bytes = vbytes.buf;
 518    }
 519
 520    if (lo < 0)
 521        lo = 0;
 522    if (hi < lo)
 523        hi = lo;
 524    if (hi > Py_SIZE(self))
 525        hi = Py_SIZE(self);
 526
 527    avail = hi - lo;
 528    if (avail < 0)
 529        lo = hi = avail = 0;
 530
 531    if (avail != needed) {
 532        if (avail > needed) {
 533            if (!_canresize(self)) {
 534                res = -1;
 535                goto finish;
 536            }
 537            /*
 538              0   lo               hi               old_size
 539              |   |<----avail----->|<-----tomove------>|
 540              |   |<-needed->|<-----tomove------>|
 541              0   lo      new_hi              new_size
 542            */
 543            memmove(self->ob_bytes + lo + needed, self->ob_bytes + hi,
 544                    Py_SIZE(self) - hi);
 545        }
 546        /* XXX(nnorwitz): need to verify this can't overflow! */
 547        if (PyByteArray_Resize((PyObject *)self,
 548                           Py_SIZE(self) + needed - avail) < 0) {
 549                res = -1;
 550                goto finish;
 551        }
 552        if (avail < needed) {
 553            /*
 554              0   lo        hi               old_size
 555              |   |<-avail->|<-----tomove------>|
 556              |   |<----needed---->|<-----tomove------>|
 557              0   lo            new_hi              new_size
 558             */
 559            memmove(self->ob_bytes + lo + needed, self->ob_bytes + hi,
 560                    Py_SIZE(self) - lo - needed);
 561        }
 562    }
 563
 564    if (needed > 0)
 565        memcpy(self->ob_bytes + lo, bytes, needed);
 566
 567
 568 finish:
 569    if (vbytes.len != -1)
 570            PyBuffer_Release(&vbytes);
 571    return res;
 572}
 573
 574static int
 575bytes_setitem(PyByteArrayObject *self, Py_ssize_t i, PyObject *value)
 576{
 577    int ival;
 578
 579    if (i < 0)
 580        i += Py_SIZE(self);
 581
 582    if (i < 0 || i >= Py_SIZE(self)) {
 583        PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
 584        return -1;
 585    }
 586
 587    if (value == NULL)
 588        return bytes_setslice(self, i, i+1, NULL);
 589
 590    if (!_getbytevalue(value, &ival))
 591        return -1;
 592
 593    self->ob_bytes[i] = ival;
 594    return 0;
 595}
 596
 597static int
 598bytes_ass_subscript(PyByteArrayObject *self, PyObject *index, PyObject *values)
 599{
 600    Py_ssize_t start, stop, step, slicelen, needed;
 601    char *bytes;
 602
 603    if (PyIndex_Check(index)) {
 604        Py_ssize_t i = PyNumber_AsSsize_t(index, PyExc_IndexError);
 605
 606        if (i == -1 && PyErr_Occurred())
 607            return -1;
 608
 609        if (i < 0)
 610            i += PyByteArray_GET_SIZE(self);
 611
 612        if (i < 0 || i >= Py_SIZE(self)) {
 613            PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
 614            return -1;
 615        }
 616
 617        if (values == NULL) {
 618            /* Fall through to slice assignment */
 619            start = i;
 620            stop = i + 1;
 621            step = 1;
 622            slicelen = 1;
 623        }
 624        else {
 625            int ival;
 626            if (!_getbytevalue(values, &ival))
 627                return -1;
 628            self->ob_bytes[i] = (char)ival;
 629            return 0;
 630        }
 631    }
 632    else if (PySlice_Check(index)) {
 633        if (PySlice_GetIndicesEx((PySliceObject *)index,
 634                                 PyByteArray_GET_SIZE(self),
 635                                 &start, &stop, &step, &slicelen) < 0) {
 636            return -1;
 637        }
 638    }
 639    else {
 640        PyErr_SetString(PyExc_TypeError, "bytearray indices must be integer");
 641        return -1;
 642    }
 643
 644    if (values == NULL) {
 645        bytes = NULL;
 646        needed = 0;
 647    }
 648    else if (values == (PyObject *)self || !PyByteArray_Check(values)) {
 649        /* Make a copy an call this function recursively */
 650        int err;
 651        values = PyByteArray_FromObject(values);
 652        if (values == NULL)
 653            return -1;
 654        err = bytes_ass_subscript(self, index, values);
 655        Py_DECREF(values);
 656        return err;
 657    }
 658    else {
 659        assert(PyByteArray_Check(values));
 660        bytes = ((PyByteArrayObject *)values)->ob_bytes;
 661        needed = Py_SIZE(values);
 662    }
 663    /* Make sure b[5:2] = ... inserts before 5, not before 2. */
 664    if ((step < 0 && start < stop) ||
 665        (step > 0 && start > stop))
 666        stop = start;
 667    if (step == 1) {
 668        if (slicelen != needed) {
 669            if (!_canresize(self))
 670                return -1;
 671            if (slicelen > needed) {
 672                /*
 673                  0   start           stop              old_size
 674                  |   |<---slicelen--->|<-----tomove------>|
 675                  |   |<-needed->|<-----tomove------>|
 676                  0   lo      new_hi              new_size
 677                */
 678                memmove(self->ob_bytes + start + needed, self->ob_bytes + stop,
 679                        Py_SIZE(self) - stop);
 680            }
 681            if (PyByteArray_Resize((PyObject *)self,
 682                               Py_SIZE(self) + needed - slicelen) < 0)
 683                return -1;
 684            if (slicelen < needed) {
 685                /*
 686                  0   lo        hi               old_size
 687                  |   |<-avail->|<-----tomove------>|
 688                  |   |<----needed---->|<-----tomove------>|
 689                  0   lo            new_hi              new_size
 690                 */
 691                memmove(self->ob_bytes + start + needed, self->ob_bytes + stop,
 692                        Py_SIZE(self) - start - needed);
 693            }
 694        }
 695
 696        if (needed > 0)
 697            memcpy(self->ob_bytes + start, bytes, needed);
 698
 699        return 0;
 700    }
 701    else {
 702        if (needed == 0) {
 703            /* Delete slice */
 704            Py_ssize_t cur, i;
 705
 706            if (!_canresize(self))
 707                return -1;
 708            if (step < 0) {
 709                stop = start + 1;
 710                start = stop + step * (slicelen - 1) - 1;
 711                step = -step;
 712            }
 713            for (cur = start, i = 0;
 714                 i < slicelen; cur += step, i++) {
 715                Py_ssize_t lim = step - 1;
 716
 717                if (cur + step >= PyByteArray_GET_SIZE(self))
 718                    lim = PyByteArray_GET_SIZE(self) - cur - 1;
 719
 720                memmove(self->ob_bytes + cur - i,
 721                        self->ob_bytes + cur + 1, lim);
 722            }
 723            /* Move the tail of the bytes, in one chunk */
 724            cur = start + slicelen*step;
 725            if (cur < PyByteArray_GET_SIZE(self)) {
 726                memmove(self->ob_bytes + cur - slicelen,
 727                        self->ob_bytes + cur,
 728                        PyByteArray_GET_SIZE(self) - cur);
 729            }
 730            if (PyByteArray_Resize((PyObject *)self,
 731                               PyByteArray_GET_SIZE(self) - slicelen) < 0)
 732                return -1;
 733
 734            return 0;
 735        }
 736        else {
 737            /* Assign slice */
 738            Py_ssize_t cur, i;
 739
 740            if (needed != slicelen) {
 741                PyErr_Format(PyExc_ValueError,
 742                             "attempt to assign bytes of size %zd "
 743                             "to extended slice of size %zd",
 744                             needed, slicelen);
 745                return -1;
 746            }
 747            for (cur = start, i = 0; i < slicelen; cur += step, i++)
 748                self->ob_bytes[cur] = bytes[i];
 749            return 0;
 750        }
 751    }
 752}
 753
 754static int
 755bytes_init(PyByteArrayObject *self, PyObject *args, PyObject *kwds)
 756{
 757    static char *kwlist[] = {"source", "encoding", "errors", 0};
 758    PyObject *arg = NULL;
 759    const char *encoding = NULL;
 760    const char *errors = NULL;
 761    Py_ssize_t count;
 762    PyObject *it;
 763    PyObject *(*iternext)(PyObject *);
 764
 765    if (Py_SIZE(self) != 0) {
 766        /* Empty previous contents (yes, do this first of all!) */
 767        if (PyByteArray_Resize((PyObject *)self, 0) < 0)
 768            return -1;
 769    }
 770
 771    /* Parse arguments */
 772    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|Oss:bytearray", kwlist,
 773                                     &arg, &encoding, &errors))
 774        return -1;
 775
 776    /* Make a quick exit if no first argument */
 777    if (arg == NULL) {
 778        if (encoding != NULL || errors != NULL) {
 779            PyErr_SetString(PyExc_TypeError,
 780                            "encoding or errors without sequence argument");
 781            return -1;
 782        }
 783        return 0;
 784    }
 785
 786    if (PyBytes_Check(arg)) {
 787        PyObject *new, *encoded;
 788        if (encoding != NULL) {
 789            encoded = PyCodec_Encode(arg, encoding, errors);
 790            if (encoded == NULL)
 791                return -1;
 792            assert(PyBytes_Check(encoded));
 793        }
 794        else {
 795            encoded = arg;
 796            Py_INCREF(arg);
 797        }
 798        new = bytes_iconcat(self, arg);
 799        Py_DECREF(encoded);
 800        if (new == NULL)
 801            return -1;
 802        Py_DECREF(new);
 803        return 0;
 804    }
 805
 806    if (PyUnicode_Check(arg)) {
 807        /* Encode via the codec registry */
 808        PyObject *encoded, *new;
 809        if (encoding == NULL) {
 810            PyErr_SetString(PyExc_TypeError,
 811                            "unicode argument without an encoding");
 812            return -1;
 813        }
 814        encoded = PyCodec_Encode(arg, encoding, errors);
 815        if (encoded == NULL)
 816            return -1;
 817        assert(PyBytes_Check(encoded));
 818        new = bytes_iconcat(self, encoded);
 819        Py_DECREF(encoded);
 820        if (new == NULL)
 821            return -1;
 822        Py_DECREF(new);
 823        return 0;
 824    }
 825
 826    /* If it's not unicode, there can't be encoding or errors */
 827    if (encoding != NULL || errors != NULL) {
 828        PyErr_SetString(PyExc_TypeError,
 829                        "encoding or errors without a string argument");
 830        return -1;
 831    }
 832
 833    /* Is it an int? */
 834    count = PyNumber_AsSsize_t(arg, PyExc_ValueError);
 835    if (count == -1 && PyErr_Occurred())
 836        PyErr_Clear();
 837    else {
 838        if (count < 0) {
 839            PyErr_SetString(PyExc_ValueError, "negative count");
 840            return -1;
 841        }
 842        if (count > 0) {
 843            if (PyByteArray_Resize((PyObject *)self, count))
 844                return -1;
 845            memset(self->ob_bytes, 0, count);
 846        }
 847        return 0;
 848    }
 849
 850    /* Use the buffer API */
 851    if (PyObject_CheckBuffer(arg)) {
 852        Py_ssize_t size;
 853        Py_buffer view;
 854        if (PyObject_GetBuffer(arg, &view, PyBUF_FULL_RO) < 0)
 855            return -1;
 856        size = view.len;
 857        if (PyByteArray_Resize((PyObject *)self, size) < 0) goto fail;
 858        if (PyBuffer_ToContiguous(self->ob_bytes, &view, size, 'C') < 0)
 859                goto fail;
 860        PyBuffer_Release(&view);
 861        return 0;
 862    fail:
 863        PyBuffer_Release(&view);
 864        return -1;
 865    }
 866
 867    /* XXX Optimize this if the arguments is a list, tuple */
 868
 869    /* Get the iterator */
 870    it = PyObject_GetIter(arg);
 871    if (it == NULL)
 872        return -1;
 873    iternext = *Py_TYPE(it)->tp_iternext;
 874
 875    /* Run the iterator to exhaustion */
 876    for (;;) {
 877        PyObject *item;
 878        int rc, value;
 879
 880        /* Get the next item */
 881        item = iternext(it);
 882        if (item == NULL) {
 883            if (PyErr_Occurred()) {
 884                if (!PyErr_ExceptionMatches(PyExc_StopIteration))
 885                    goto error;
 886                PyErr_Clear();
 887            }
 888            break;
 889        }
 890
 891        /* Interpret it as an int (__index__) */
 892        rc = _getbytevalue(item, &value);
 893        Py_DECREF(item);
 894        if (!rc)
 895            goto error;
 896
 897        /* Append the byte */
 898        if (Py_SIZE(self) < self->ob_alloc)
 899            Py_SIZE(self)++;
 900        else if (PyByteArray_Resize((PyObject *)self, Py_SIZE(self)+1) < 0)
 901            goto error;
 902        self->ob_bytes[Py_SIZE(self)-1] = value;
 903    }
 904
 905    /* Clean up and return success */
 906    Py_DECREF(it);
 907    return 0;
 908
 909 error:
 910    /* Error handling when it != NULL */
 911    Py_DECREF(it);
 912    return -1;
 913}
 914
 915/* Mostly copied from string_repr, but without the
 916   "smart quote" functionality. */
 917static PyObject *
 918bytes_repr(PyByteArrayObject *self)
 919{
 920    static const char *hexdigits = "0123456789abcdef";
 921    const char *quote_prefix = "bytearray(b";
 922    const char *quote_postfix = ")";
 923    Py_ssize_t length = Py_SIZE(self);
 924    /* 14 == strlen(quote_prefix) + 2 + strlen(quote_postfix) */
 925    size_t newsize = 14 + 4 * length;
 926    PyObject *v;
 927    if (newsize > PY_SSIZE_T_MAX || newsize / 4 - 3 != length) {
 928        PyErr_SetString(PyExc_OverflowError,
 929            "bytearray object is too large to make repr");
 930        return NULL;
 931    }
 932    v = PyUnicode_FromUnicode(NULL, newsize);
 933    if (v == NULL) {
 934        return NULL;
 935    }
 936    else {
 937        register Py_ssize_t i;
 938        register Py_UNICODE c;
 939        register Py_UNICODE *p;
 940        int quote;
 941
 942        /* Figure out which quote to use; single is preferred */
 943        quote = '\'';
 944        {
 945            char *test, *start;
 946            start = PyByteArray_AS_STRING(self);
 947            for (test = start; test < start+length; ++test) {
 948                if (*test == '"') {
 949                    quote = '\''; /* back to single */
 950                    goto decided;
 951                }
 952                else if (*test == '\'')
 953                    quote = '"';
 954            }
 955          decided:
 956            ;
 957        }
 958
 959        p = PyUnicode_AS_UNICODE(v);
 960        while (*quote_prefix)
 961            *p++ = *quote_prefix++;
 962        *p++ = quote;
 963
 964        for (i = 0; i < length; i++) {
 965            /* There's at least enough room for a hex escape
 966               and a closing quote. */
 967            assert(newsize - (p - PyUnicode_AS_UNICODE(v)) >= 5);
 968            c = self->ob_bytes[i];
 969            if (c == '\'' || c == '\\')
 970                *p++ = '\\', *p++ = c;
 971            else if (c == '\t')
 972                *p++ = '\\', *p++ = 't';
 973            else if (c == '\n')
 974                *p++ = '\\', *p++ = 'n';
 975            else if (c == '\r')
 976                *p++ = '\\', *p++ = 'r';
 977            else if (c == 0)
 978                *p++ = '\\', *p++ = 'x', *p++ = '0', *p++ = '0';
 979            else if (c < ' ' || c >= 0x7f) {
 980                *p++ = '\\';
 981                *p++ = 'x';
 982                *p++ = hexdigits[(c & 0xf0) >> 4];
 983                *p++ = hexdigits[c & 0xf];
 984            }
 985            else
 986                *p++ = c;
 987        }
 988        assert(newsize - (p - PyUnicode_AS_UNICODE(v)) >= 1);
 989        *p++ = quote;
 990        while (*quote_postfix) {
 991           *p++ = *quote_postfix++;
 992        }
 993        *p = '\0';
 994        if (PyUnicode_Resize(&v, (p - PyUnicode_AS_UNICODE(v)))) {
 995            Py_DECREF(v);
 996            return NULL;
 997        }
 998        return v;
 999    }
1000}
1001
1002static PyObject *
1003bytes_str(PyObject *op)
1004{
1005#if 0
1006    if (Py_BytesWarningFlag) {
1007        if (PyErr_WarnEx(PyExc_BytesWarning,
1008                 "str() on a bytearray instance", 1))
1009            return NULL;
1010    }
1011    return bytes_repr((PyByteArrayObject*)op);
1012#endif
1013    return PyBytes_FromStringAndSize(((PyByteArrayObject*)op)->ob_bytes, Py_SIZE(op));
1014}
1015
1016static PyObject *
1017bytes_richcompare(PyObject *self, PyObject *other, int op)
1018{
1019    Py_ssize_t self_size, other_size;
1020    Py_buffer self_bytes, other_bytes;
1021    PyObject *res;
1022    Py_ssize_t minsize;
1023    int cmp;
1024
1025    /* Bytes can be compared to anything that supports the (binary)
1026       buffer API.  Except that a comparison with Unicode is always an
1027       error, even if the comparison is for equality. */
1028    if (PyObject_IsInstance(self, (PyObject*)&PyUnicode_Type) ||
1029        PyObject_IsInstance(other, (PyObject*)&PyUnicode_Type)) {
1030        if (Py_BytesWarningFlag && op == Py_EQ) {
1031            if (PyErr_WarnEx(PyExc_BytesWarning,
1032                            "Comparsion between bytearray and string", 1))
1033                return NULL;
1034        }
1035
1036        Py_INCREF(Py_NotImplemented);
1037        return Py_NotImplemented;
1038    }
1039
1040    self_size = _getbuffer(self, &self_bytes);
1041    if (self_size < 0) {
1042        PyErr_Clear();
1043        Py_INCREF(Py_NotImplemented);
1044        return Py_NotImplemented;
1045    }
1046
1047    other_size = _getbuffer(other, &other_bytes);
1048    if (other_size < 0) {
1049        PyErr_Clear();
1050        PyBuffer_Release(&self_bytes);
1051        Py_INCREF(Py_NotImplemented);
1052        return Py_NotImplemented;
1053    }
1054
1055    if (self_size != other_size && (op == Py_EQ || op == Py_NE)) {
1056        /* Shortcut: if the lengths differ, the objects differ */
1057        cmp = (op == Py_NE);
1058    }
1059    else {
1060        minsize = self_size;
1061        if (other_size < minsize)
1062            minsize = other_size;
1063
1064        cmp = memcmp(self_bytes.buf, other_bytes.buf, minsize);
1065        /* In ISO C, memcmp() guarantees to use unsigned bytes! */
1066
1067        if (cmp == 0) {
1068            if (self_size < other_size)
1069                cmp = -1;
1070            else if (self_size > other_size)
1071                cmp = 1;
1072        }
1073
1074        switch (op) {
1075        case Py_LT: cmp = cmp <  0; break;
1076        case Py_LE: cmp = cmp <= 0; break;
1077        case Py_EQ: cmp = cmp == 0; break;
1078        case Py_NE: cmp = cmp != 0; break;
1079        case Py_GT: cmp = cmp >  0; break;
1080        case Py_GE: cmp = cmp >= 0; break;
1081        }
1082    }
1083
1084    res = cmp ? Py_True : Py_False;
1085    PyBuffer_Release(&self_bytes);
1086    PyBuffer_Release(&other_bytes);
1087    Py_INCREF(res);
1088    return res;
1089}
1090
1091static void
1092bytes_dealloc(PyByteArrayObject *self)
1093{
1094	if (self->ob_exports > 0) {
1095		PyErr_SetString(PyExc_SystemError,
1096                        "deallocated bytearray object has exported buffers");
1097		PyErr_Print();
1098	}
1099    if (self->ob_bytes != 0) {
1100        PyMem_Free(self->ob_bytes);
1101    }
1102    Py_TYPE(self)->tp_free((PyObject *)self);
1103}
1104
1105
1106/* -------------------------------------------------------------------- */
1107/* Methods */
1108
1109#define STRINGLIB_CHAR char
1110#define STRINGLIB_CMP memcmp
1111#define STRINGLIB_LEN PyByteArray_GET_SIZE
1112#define STRINGLIB_STR PyByteArray_AS_STRING
1113#define STRINGLIB_NEW PyByteArray_FromStringAndSize
1114#define STRINGLIB_EMPTY nullbytes
1115#define STRINGLIB_CHECK_EXACT PyByteArray_CheckExact
1116#define STRINGLIB_MUTABLE 1
1117#define FROM_BYTEARRAY 1
1118
1119#include "stringlib/fastsearch.h"
1120#include "stringlib/count.h"
1121#include "stringlib/find.h"
1122#include "stringlib/partition.h"
1123#include "stringlib/ctype.h"
1124#include "stringlib/transmogrify.h"
1125
1126
1127/* The following Py_LOCAL_INLINE and Py_LOCAL functions
1128were copied from the old char* style string object. */
1129
1130Py_LOCAL_INLINE(void)
1131_adjust_indices(Py_ssize_t *start, Py_ssize_t *end, Py_ssize_t len)
1132{
1133    if (*end > len)
1134        *end = len;
1135    else if (*end < 0)
1136        *end += len;
1137    if (*end < 0)
1138        *end = 0;
1139    if (*start < 0)
1140        *start += len;
1141    if (*start < 0)
1142        *start = 0;
1143}
1144
1145
1146Py_LOCAL_INLINE(Py_ssize_t)
1147bytes_find_internal(PyByteArrayObject *self, PyObject *args, int dir)
1148{
1149    PyObject *subobj;
1150    Py_buffer subbuf;
1151    Py_ssize_t start=0, end=PY_SSIZE_T_MAX;
1152    Py_ssize_t res;
1153
1154    if (!PyArg_ParseTuple(args, "O|O&O&:find/rfind/index/rindex", &subobj,
1155        _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1156        return -2;
1157    if (_getbuffer(subobj, &subbuf) < 0)
1158        return -2;
1159    if (dir > 0)
1160        res = stringlib_find_slice(
1161            PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
1162            subbuf.buf, subbuf.len, start, end);
1163    else
1164        res = stringlib_rfind_slice(
1165            PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
1166            subbuf.buf, subbuf.len, start, end);
1167    PyBuffer_Release(&subbuf);
1168    return res;
1169}
1170
1171PyDoc_STRVAR(find__doc__,
1172"B.find(sub [,start [,end]]) -> int\n\
1173\n\
1174Return the lowest index in B where subsection sub is found,\n\
1175such that sub is contained within s[start,end].  Optional\n\
1176arguments start and end are interpreted as in slice notation.\n\
1177\n\
1178Return -1 on failure.");
1179
1180static PyObject *
1181bytes_find(PyByteArrayObject *self, PyObject *args)
1182{
1183    Py_ssize_t result = bytes_find_internal(self, args, +1);
1184    if (result == -2)
1185        return NULL;
1186    return PyInt_FromSsize_t(result);
1187}
1188
1189PyDoc_STRVAR(count__doc__,
1190"B.count(sub [,start [,end]]) -> int\n\
1191\n\
1192Return the number of non-overlapping occurrences of subsection sub in\n\
1193bytes B[start:end].  Optional arguments start and end are interpreted\n\
1194as in slice notation.");
1195
1196static PyObject *
1197bytes_count(PyByteArrayObject *self, PyObject *args)
1198{
1199    PyObject *sub_obj;
1200    const char *str = PyByteArray_AS_STRING(self);
1201    Py_ssize_t start = 0, end = PY_SSIZE_T_MAX;
1202    Py_buffer vsub;
1203    PyObject *count_obj;
1204
1205    if (!PyArg_ParseTuple(args, "O|O&O&:count", &sub_obj,
1206        _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1207        return NULL;
1208
1209    if (_getbuffer(sub_obj, &vsub) < 0)
1210        return NULL;
1211
1212    _adjust_indices(&start, &end, PyByteArray_GET_SIZE(self));
1213
1214    count_obj = PyInt_FromSsize_t(
1215        stringlib_count(str + start, end - start, vsub.buf, vsub.len)
1216        );
1217    PyBuffer_Release(&vsub);
1218    return count_obj;
1219}
1220
1221
1222PyDoc_STRVAR(index__doc__,
1223"B.index(sub [,start [,end]]) -> int\n\
1224\n\
1225Like B.find() but raise ValueError when the subsection is not found.");
1226
1227static PyObject *
1228bytes_index(PyByteArrayObject *self, PyObject *args)
1229{
1230    Py_ssize_t result = bytes_find_internal(self, args, +1);
1231    if (result == -2)
1232        return NULL;
1233    if (result == -1) {
1234        PyErr_SetString(PyExc_ValueError,
1235                        "subsection not found");
1236        return NULL;
1237    }
1238    return PyInt_FromSsize_t(result);
1239}
1240
1241
1242PyDoc_STRVAR(rfind__doc__,
1243"B.rfind(sub [,start [,end]]) -> int\n\
1244\n\
1245Return the highest index in B where subsection sub is found,\n\
1246such that sub is contained within s[start,end].  Optional\n\
1247arguments start and end are interpreted as in slice notation.\n\
1248\n\
1249Return -1 on failure.");
1250
1251static PyObject *
1252bytes_rfind(PyByteArrayObject *self, PyObject *args)
1253{
1254    Py_ssize_t result = bytes_find_internal(self, args, -1);
1255    if (result == -2)
1256        return NULL;
1257    return PyInt_FromSsize_t(result);
1258}
1259
1260
1261PyDoc_STRVAR(rindex__doc__,
1262"B.rindex(sub [,start [,end]]) -> int\n\
1263\n\
1264Like B.rfind() but raise ValueError when the subsection is not found.");
1265
1266static PyObject *
1267bytes_rindex(PyByteArrayObject *self, PyObject *args)
1268{
1269    Py_ssize_t result = bytes_find_internal(self, args, -1);
1270    if (result == -2)
1271        return NULL;
1272    if (result == -1) {
1273        PyErr_SetString(PyExc_ValueError,
1274                        "subsection not found");
1275        return NULL;
1276    }
1277    return PyInt_FromSsize_t(result);
1278}
1279
1280
1281static int
1282bytes_contains(PyObject *self, PyObject *arg)
1283{
1284    Py_ssize_t ival = PyNumber_AsSsize_t(arg, PyExc_ValueError);
1285    if (ival == -1 && PyErr_Occurred()) {
1286        Py_buffer varg;
1287        int pos;
1288        PyErr_Clear();
1289        if (_getbuffer(arg, &varg) < 0)
1290            return -1;
1291        pos = stringlib_find(PyByteArray_AS_STRING(self), Py_SIZE(self),
1292                             varg.buf, varg.len, 0);
1293        PyBuffer_Release(&varg);
1294        return pos >= 0;
1295    }
1296    if (ival < 0 || ival >= 256) {
1297        PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
1298        return -1;
1299    }
1300
1301    return memchr(PyByteArray_AS_STRING(self), ival, Py_SIZE(self)) != NULL;
1302}
1303
1304
1305/* Matches the end (direction >= 0) or start (direction < 0) of self
1306 * against substr, using the start and end arguments. Returns
1307 * -1 on error, 0 if not found and 1 if found.
1308 */
1309Py_LOCAL(int)
1310_bytes_tailmatch(PyByteArrayObject *self, PyObject *substr, Py_ssize_t start,
1311                 Py_ssize_t end, int direction)
1312{
1313    Py_ssize_t len = PyByteArray_GET_SIZE(self);
1314    const char* str;
1315    Py_buffer vsubstr;
1316    int rv = 0;
1317
1318    str = PyByteArray_AS_STRING(self);
1319
1320    if (_getbuffer(substr, &vsubstr) < 0)
1321        return -1;
1322
1323    _adjust_indices(&start, &end, len);
1324
1325    if (direction < 0) {
1326        /* startswith */
1327        if (start+vsubstr.len > len) {
1328            goto done;
1329        }
1330    } else {
1331        /* endswith */
1332        if (end-start < vsubstr.len || start > len) {
1333            goto done;
1334        }
1335
1336        if (end-vsubstr.len > start)
1337            start = end - vsubstr.len;
1338    }
1339    if (end-start >= vsubstr.len)
1340        rv = ! memcmp(str+start, vsubstr.buf, vsubstr.len);
1341
1342done:
1343    PyBuffer_Release(&vsubstr);
1344    return rv;
1345}
1346
1347
1348PyDoc_STRVAR(startswith__doc__,
1349"B.startswith(prefix [,start [,end]]) -> bool\n\
1350\n\
1351Return True if B starts with the specified prefix, False otherwise.\n\
1352With optional start, test B beginning at that position.\n\
1353With optional end, stop comparing B at that position.\n\
1354prefix can also be a tuple of strings to try.");
1355
1356static PyObject *
1357bytes_startswith(PyByteArrayObject *self, PyObject *args)
1358{
1359    Py_ssize_t start = 0;
1360    Py_ssize_t end = PY_SSIZE_T_MAX;
1361    PyObject *subobj;
1362    int result;
1363
1364    if (!PyArg_ParseTuple(args, "O|O&O&:startswith", &subobj,
1365        _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1366        return NULL;
1367    if (PyTuple_Check(subobj)) {
1368        Py_ssize_t i;
1369        for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
1370            result = _bytes_tailmatch(self,
1371                                      PyTuple_GET_ITEM(subobj, i),
1372                                      start, end, -1);
1373            if (result == -1)
1374                return NULL;
1375            else if (result) {
1376                Py_RETURN_TRUE;
1377            }
1378        }
1379        Py_RETURN_FALSE;
1380    }
1381    result = _bytes_tailmatch(self, subobj, start, end, -1);
1382    if (result == -1)
1383        return NULL;
1384    else
1385        return PyBool_FromLong(result);
1386}
1387
1388PyDoc_STRVAR(endswith__doc__,
1389"B.endswith(suffix [,start [,end]]) -> bool\n\
1390\n\
1391Return True if B ends with the specified suffix, False otherwise.\n\
1392With optional start, test B beginning at that position.\n\
1393With optional end, stop comparing B at that position.\n\
1394suffix can also be a tuple of strings to try.");
1395
1396static PyObject *
1397bytes_endswith(PyByteArrayObject *self, PyObject *args)
1398{
1399    Py_ssize_t start = 0;
1400    Py_ssize_t end = PY_SSIZE_T_MAX;
1401    PyObject *subobj;
1402    int result;
1403
1404    if (!PyArg_ParseTuple(args, "O|O&O&:endswith", &subobj,
1405        _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1406        return NULL;
1407    if (PyTuple_Check(subobj)) {
1408        Py_ssize_t i;
1409        for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
1410            result = _bytes_tailmatch(self,
1411                                      PyTuple_GET_ITEM(subobj, i),
1412                                      start, end, +1);
1413            if (result == -1)
1414                return NULL;
1415            else if (result) {
1416                Py_RETURN_TRUE;
1417            }
1418        }
1419        Py_RETURN_FALSE;
1420    }
1421    result = _bytes_tailmatch(self, subobj, start, end, +1);
1422    if (result == -1)
1423        return NULL;
1424    else
1425        return PyBool_FromLong(result);
1426}
1427
1428
1429PyDoc_STRVAR(translate__doc__,
1430"B.translate(table[, deletechars]) -> bytearray\n\
1431\n\
1432Return a copy of B, where all characters occurring in the\n\
1433optional argument deletechars are removed, and the remaining\n\
1434characters have been mapped through the given translation\n\
1435table, which must be a bytes object of length 256.");
1436
1437static PyObject *
1438bytes_translate(PyByteArrayObject *self, PyObject *args)
1439{
1440    register char *input, *output;
1441    register const char *table;
1442    register Py_ssize_t i, c;
1443    PyObject *input_obj = (PyObject*)self;
1444    const char *output_start;
1445    Py_ssize_t inlen;
1446    PyObject *result;
1447    int trans_table[256];
1448    PyObject *tableobj, *delobj = NULL;
1449    Py_buffer vtable, vdel;
1450
1451    if (!PyArg_UnpackTuple(args, "translate", 1, 2,
1452                           &tableobj, &delobj))
1453          return NULL;
1454
1455    if (_getbuffer(tableobj, &vtable) < 0)
1456        return NULL;
1457
1458    if (vtable.len != 256) {
1459        PyErr_SetString(PyExc_ValueError,
1460                        "translation table must be 256 characters long");
1461        PyBuffer_Release(&vtable);
1462        return NULL;
1463    }
1464
1465    if (delobj != NULL) {
1466        if (_getbuffer(delobj, &vdel) < 0) {
1467            PyBuffer_Release(&vtable);
1468	    return NULL;
1469        }
1470    }
1471    else {
1472        vdel.buf = NULL;
1473        vdel.len = 0;
1474    }
1475
1476    table = (const char *)vtable.buf;
1477    inlen = PyByteArray_GET_SIZE(input_obj);
1478    result = PyByteArray_FromStringAndSize((char *)NULL, inlen);
1479    if (result == NULL)
1480        goto done;
1481    output_start = output = PyByteArray_AsString(result);
1482    input = PyByteArray_AS_STRING(input_obj);
1483
1484    if (vdel.len == 0) {
1485        /* If no deletions are required, use faster code */
1486        for (i = inlen; --i >= 0; ) {
1487            c = Py_CHARMASK(*input++);
1488            *output++ = table[c];
1489        }
1490        goto done;
1491    }
1492    
1493    for (i = 0; i < 256; i++)
1494        trans_table[i] = Py_CHARMASK(table[i]);
1495
1496    for (i = 0; i < vdel.len; i++)
1497        trans_table[(int) Py_CHARMASK( ((unsigned char*)vdel.buf)[i] )] = -1;
1498
1499    for (i = inlen; --i >= 0; ) {
1500        c = Py_CHARMASK(*input++);
1501        if (trans_table[c] != -1)
1502            if (Py_CHARMASK(*output++ = (char)trans_table[c]) == c)
1503                    continue;
1504    }
1505    /* Fix the size of the resulting string */
1506    if (inlen > 0)
1507        PyByteArray_Resize(result, output - output_start);
1508
1509done:
1510    PyBuffer_Release(&vtable);
1511    if (delobj != NULL)
1512        PyBuffer_Release(&vdel);
1513    return result;
1514}
1515
1516
1517#define FORWARD 1
1518#define REVERSE -1
1519
1520/* find and count characters and substrings */
1521
1522#define findchar(target, target_len, c)                         \
1523  ((char *)memchr((const void *)(target), c, target_len))
1524
1525/* Don't call if length < 2 */
1526#define Py_STRING_MATCH(target, offset, pattern, length)        \
1527  (target[offset] == pattern[0] &&                              \
1528   target[offset+length-1] == pattern[length-1] &&              \
1529   !memcmp(target+offset+1, pattern+1, length-2) )
1530
1531
1532/* Bytes ops must return a string, create a copy */
1533Py_LOCAL(PyByteArrayObject *)
1534return_self(PyByteArrayObject *self)
1535{
1536    return (PyByteArrayObject *)PyByteArray_FromStringAndSize(
1537            PyByteArray_AS_STRING(self),
1538            PyByteArray_GET_SIZE(self));
1539}
1540
1541Py_LOCAL_INLINE(Py_ssize_t)
1542countchar(const char *target, Py_ssize_t target_len, char c, Py_ssize_t maxcount)
1543{
1544    Py_ssize_t count=0;
1545    const char *start=target;
1546    const char *end=target+target_len;
1547
1548    while ( (start=findchar(start, end-start, c)) != NULL ) {
1549        count++;
1550        if (count >= maxcount)
1551            break;
1552        start += 1;
1553    }
1554    return count;
1555}
1556
1557Py_LOCAL(Py_ssize_t)
1558findstring(const char *target, Py_ssize_t target_len,
1559           const char *pattern, Py_ssize_t pattern_len,
1560           Py_ssize_t start,
1561           Py_ssize_t end,
1562           int direction)
1563{
1564    if (start < 0) {
1565        start += target_len;
1566        if (start < 0)
1567            start = 0;
1568    }
1569    if (end > target_len) {
1570        end = target_len;
1571    } else if (end < 0) {
1572        end += target_len;
1573        if (end < 0)
1574            end = 0;
1575    }
1576
1577    /* zero-length substrings always match at the first attempt */
1578    if (pattern_len == 0)
1579        return (direction > 0) ? start : end;
1580
1581    end -= pattern_len;
1582
1583    if (direction < 0) {
1584        for (; end >= start; end--)
1585            if (Py_STRING_MATCH(target, end, pattern, pattern_len))
1586                return end;
1587    } else {
1588        for (; start <= end; start++)
1589            if (Py_STRING_MATCH(target, start, pattern, pattern_len))
1590                return start;
1591    }
1592    return -1;
1593}
1594
1595Py_LOCAL_INLINE(Py_ssize_t)
1596countstring(const char *target, Py_ssize_t target_len,
1597            const char *pattern, Py_ssize_t pattern_len,
1598            Py_ssize_t start,
1599            Py_ssize_t end,
1600            int direction, Py_ssize_t maxcount)
1601{
1602    Py_ssize_t count=0;
1603
1604    if (start < 0) {
1605        start += target_len;
1606        if (start < 0)
1607            start = 0;
1608    }
1609    if (end > target_len) {
1610        end = target_len;
1611    } else if (end < 0) {
1612        end += target_len;
1613        if (end < 0)
1614            end = 0;
1615    }
1616
1617    /* zero-length substrings match everywhere */
1618    if (pattern_len == 0 || maxcount == 0) {
1619        if (target_len+1 < maxcount)
1620            return target_len+1;
1621        return maxcount;
1622    }
1623
1624    end -= pattern_len;
1625    if (direction < 0) {
1626        for (; (end >= start); end--)
1627            if (Py_STRING_MATCH(target, end, pattern, pattern_len)) {
1628                count++;
1629                if (--maxcount <= 0) break;
1630                end -= pattern_len-1;
1631            }
1632    } else {
1633        for (; (start <= end); start++)
1634            if (Py_STRING_MATCH(target, start, pattern, pattern_len)) {
1635                count++;
1636                if (--maxcount <= 0)
1637                    break;
1638                start += pattern_len-1;
1639            }
1640    }
1641    return count;
1642}
1643
1644
1645/* Algorithms for different cases of string replacement */
1646
1647/* len(self)>=1, from="", len(to)>=1, maxcount>=1 */
1648Py_LOCAL(PyByteArrayObject *)
1649replace_interleave(PyByteArrayObject *self,
1650                   const char *to_s, Py_ssize_t to_len,
1651                   Py_ssize_t maxcount)
1652{
1653    char *self_s, *result_s;
1654    Py_ssize_t self_len, result_len;
1655    Py_ssize_t count, i, product;
1656    PyByteArrayObject *result;
1657
1658    self_len = PyByteArray_GET_SIZE(self);
1659
1660    /* 1 at the end plus 1 after every character */
1661    count = self_len+1;
1662    if (maxcount < count)
1663        count = maxcount;
1664
1665    /* Check for overflow */
1666    /*   result_len = count * to_len + self_len; */
1667    product = count * to_len;
1668    if (product / to_len != count) {
1669        PyErr_SetString(PyExc_OverflowError,
1670                        "replace string is too long");
1671        return NULL;
1672    }
1673    result_len = product + self_len;
1674    if (result_len < 0) {
1675        PyErr_SetString(PyExc_OverflowError,
1676                        "replace string is too long");
1677        return NULL;
1678    }
1679
1680    if (! (result = (PyByteArrayObject *)
1681                     PyByteArray_FromStringAndSize(NULL, result_len)) )
1682        return NULL;
1683
1684    self_s = PyByteArray_AS_STRING(self);
1685    result_s = PyByteArray_AS_STRING(result);
1686
1687    /* TODO: special case single character, which doesn't need memcpy */
1688
1689    /* Lay the first one down (guaranteed this will occur) */
1690    Py_MEMCPY(result_s, to_s, to_len);
1691    result_s += to_len;
1692    count -= 1;
1693
1694    for (i=0; i<count; i++) {
1695        *result_s++ = *self_s++;
1696        Py_MEMCPY(result_s, to_s, to_len);
1697        result_s += to_len;
1698    }
1699
1700    /* Copy the rest of the original string */
1701    Py_MEMCPY(result_s, self_s, self_len-i);
1702
1703    return result;
1704}
1705
1706/* Special case for deleting a single character */
1707/* len(self)>=1, len(from)==1, to="", maxcount>=1 */
1708Py_LOCAL(PyByteArrayObject *)
1709replace_delete_single_character(PyByteArrayObject *self,
1710                                char from_c, Py_ssize_t maxcount)
1711{
1712    char *self_s, *result_s;
1713    char *start, *next, *end;
1714    Py_ssize_t self_len, result_len;
1715    Py_ssize_t count;
1716    PyByteArrayObject *result;
1717
1718    self_len = PyByteArray_GET_SIZE(self);
1719    self_s = PyByteArray_AS_STRING(self);
1720
1721    count = countchar(self_s, self_len, from_c, maxcount);
1722    if (count == 0) {
1723        return return_self(self);
1724    }
1725
1726    result_len = self_len - count;  /* from_len == 1 */
1727    assert(result_len>=0);
1728
1729    if ( (result = (PyByteArrayObject *)
1730                    PyByteArray_FromStringAndSize(NULL, result_len)) == NULL)
1731        return NULL;
1732    result_s = PyByteArray_AS_STRING(result);
1733
1734    start = self_s;
1735    end = self_s + self_len;
1736    while (count-- > 0) {
1737        next = findchar(start, end-start, from_c);
1738        if (next == NULL)
1739            break;
1740        Py_MEMCPY(result_s, start, next-start);
1741        result_s += (next-start);
1742        start = next+1;
1743    }
1744    Py_MEMCPY(result_s, start, end-start);
1745
1746    return result;
1747}
1748
1749/* len(self)>=1, len(from)>=2, to="", maxcount>=1 */
1750
1751Py_LOCAL(PyByteArrayObject *)
1752replace_delete_substring(PyByteArrayObject *self,
1753                         const char *from_s, Py_ssize_t from_len,
1754                         Py_ssize_t maxcount)
1755{
1756    char *self_s, *result_s;
1757    char *start, *next, *end;
1758    Py_ssize_t self_len, result_len;
1759    Py_ssize_t count, offset;
1760    PyByteArrayObject *result;
1761
1762    self_len = PyByteArray_GET_SIZE(self);
1763    self_s = PyByteArray_AS_STRING(self);
1764
1765    count = countstring(self_s, self_len,
1766                        from_s, from_len,
1767                        0, self_len, 1,
1768                        maxcount);
1769
1770    if (count == 0) {
1771        /* no matches */
1772        return return_self(self);
1773    }
1774
1775    result_len = self_len - (count * from_len);
1776    assert (result_len>=0);
1777
1778    if ( (result = (PyByteArrayObject *)
1779        PyByteArray_FromStringAndSize(NULL, result_len)) == NULL )
1780            return NULL;
1781
1782    result_s = PyByteArray_AS_STRING(result);
1783
1784    start = self_s;
1785    end = self_s + self_len;
1786    while (count-- > 0) {
1787        offset = findstring(start, end-start,
1788                            from_s, from_len,
1789                            0, end-start, FORWARD);
1790        if (offset == -1)
1791            break;
1792        next = start + offset;
1793
1794        Py_MEMCPY(result_s, start, next-start);
1795
1796        r…
Large files files are truncated, but you can click here to view the full file