blob: 27affb53d2273134ebb2eaf0c908417db2be4548 [file] [log] [blame]
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001/* PyByteArray (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;
Antoine Pitrou20d6c152010-01-17 12:43:00 +00009char _PyByteArray_empty_string[] = "";
Christian Heimes2c9c7a52008-05-26 13:42:13 +000010
11void
12PyByteArray_Fini(void)
13{
14 Py_CLEAR(nullbytes);
15}
16
17int
18PyByteArray_Init(void)
19{
20 nullbytes = PyObject_New(PyByteArrayObject, &PyByteArray_Type);
21 if (nullbytes == NULL)
22 return 0;
23 nullbytes->ob_bytes = NULL;
24 Py_SIZE(nullbytes) = nullbytes->ob_alloc = 0;
25 nullbytes->ob_exports = 0;
26 return 1;
27}
28
29/* end nullbytes support */
30
31/* Helpers */
32
33static int
34_getbytevalue(PyObject* arg, int *value)
35{
36 long face_value;
37
38 if (PyLong_Check(arg)) {
39 face_value = PyLong_AsLong(arg);
Georg Brandl9a54d7c2008-07-16 23:15:30 +000040 } else {
41 PyObject *index = PyNumber_Index(arg);
42 if (index == NULL) {
43 PyErr_Format(PyExc_TypeError, "an integer is required");
Christian Heimes2c9c7a52008-05-26 13:42:13 +000044 return 0;
45 }
Georg Brandl9a54d7c2008-07-16 23:15:30 +000046 face_value = PyLong_AsLong(index);
47 Py_DECREF(index);
48 }
49
50 if (face_value < 0 || face_value >= 256) {
51 /* this includes the OverflowError in case the long is too large */
52 PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
Christian Heimes2c9c7a52008-05-26 13:42:13 +000053 return 0;
54 }
55
56 *value = face_value;
57 return 1;
58}
59
60static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +000061bytearray_getbuffer(PyByteArrayObject *obj, Py_buffer *view, int flags)
Christian Heimes2c9c7a52008-05-26 13:42:13 +000062{
63 int ret;
64 void *ptr;
65 if (view == NULL) {
66 obj->ob_exports++;
67 return 0;
68 }
Antoine Pitrou20d6c152010-01-17 12:43:00 +000069 ptr = (void *) PyByteArray_AS_STRING(obj);
Martin v. Löwis423be952008-08-13 15:53:07 +000070 ret = PyBuffer_FillInfo(view, (PyObject*)obj, ptr, Py_SIZE(obj), 0, flags);
Christian Heimes2c9c7a52008-05-26 13:42:13 +000071 if (ret >= 0) {
72 obj->ob_exports++;
73 }
74 return ret;
75}
76
77static void
Benjamin Peterson153c70f2009-04-18 15:42:12 +000078bytearray_releasebuffer(PyByteArrayObject *obj, Py_buffer *view)
Christian Heimes2c9c7a52008-05-26 13:42:13 +000079{
80 obj->ob_exports--;
81}
82
83static Py_ssize_t
84_getbuffer(PyObject *obj, Py_buffer *view)
85{
86 PyBufferProcs *buffer = Py_TYPE(obj)->tp_as_buffer;
87
88 if (buffer == NULL || buffer->bf_getbuffer == NULL)
89 {
90 PyErr_Format(PyExc_TypeError,
91 "Type %.100s doesn't support the buffer API",
92 Py_TYPE(obj)->tp_name);
93 return -1;
94 }
95
96 if (buffer->bf_getbuffer(obj, view, PyBUF_SIMPLE) < 0)
97 return -1;
98 return view->len;
99}
100
Antoine Pitrou5504e892008-12-06 21:27:53 +0000101static int
102_canresize(PyByteArrayObject *self)
103{
104 if (self->ob_exports > 0) {
105 PyErr_SetString(PyExc_BufferError,
106 "Existing exports of data: object cannot be re-sized");
107 return 0;
108 }
109 return 1;
110}
111
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000112/* Direct API functions */
113
114PyObject *
115PyByteArray_FromObject(PyObject *input)
116{
117 return PyObject_CallFunctionObjArgs((PyObject *)&PyByteArray_Type,
118 input, NULL);
119}
120
121PyObject *
122PyByteArray_FromStringAndSize(const char *bytes, Py_ssize_t size)
123{
124 PyByteArrayObject *new;
125 Py_ssize_t alloc;
126
127 if (size < 0) {
128 PyErr_SetString(PyExc_SystemError,
129 "Negative size passed to PyByteArray_FromStringAndSize");
130 return NULL;
131 }
132
Neal Norwitz3ce5d922008-08-24 07:08:55 +0000133 /* Prevent buffer overflow when setting alloc to size+1. */
134 if (size == PY_SSIZE_T_MAX) {
135 return PyErr_NoMemory();
136 }
137
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000138 new = PyObject_New(PyByteArrayObject, &PyByteArray_Type);
139 if (new == NULL)
140 return NULL;
141
142 if (size == 0) {
143 new->ob_bytes = NULL;
144 alloc = 0;
145 }
146 else {
147 alloc = size + 1;
148 new->ob_bytes = PyMem_Malloc(alloc);
149 if (new->ob_bytes == NULL) {
150 Py_DECREF(new);
151 return PyErr_NoMemory();
152 }
Antoine Pitrou20d6c152010-01-17 12:43:00 +0000153 if (bytes != NULL && size > 0)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000154 memcpy(new->ob_bytes, bytes, size);
155 new->ob_bytes[size] = '\0'; /* Trailing null byte */
156 }
157 Py_SIZE(new) = size;
158 new->ob_alloc = alloc;
159 new->ob_exports = 0;
160
161 return (PyObject *)new;
162}
163
164Py_ssize_t
165PyByteArray_Size(PyObject *self)
166{
167 assert(self != NULL);
168 assert(PyByteArray_Check(self));
169
170 return PyByteArray_GET_SIZE(self);
171}
172
173char *
174PyByteArray_AsString(PyObject *self)
175{
176 assert(self != NULL);
177 assert(PyByteArray_Check(self));
178
179 return PyByteArray_AS_STRING(self);
180}
181
182int
183PyByteArray_Resize(PyObject *self, Py_ssize_t size)
184{
185 void *sval;
186 Py_ssize_t alloc = ((PyByteArrayObject *)self)->ob_alloc;
187
188 assert(self != NULL);
189 assert(PyByteArray_Check(self));
190 assert(size >= 0);
191
Antoine Pitrou5504e892008-12-06 21:27:53 +0000192 if (size == Py_SIZE(self)) {
193 return 0;
194 }
195 if (!_canresize((PyByteArrayObject *)self)) {
196 return -1;
197 }
198
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000199 if (size < alloc / 2) {
200 /* Major downsize; resize down to exact size */
201 alloc = size + 1;
202 }
203 else if (size < alloc) {
204 /* Within allocated size; quick exit */
205 Py_SIZE(self) = size;
206 ((PyByteArrayObject *)self)->ob_bytes[size] = '\0'; /* Trailing null */
207 return 0;
208 }
209 else if (size <= alloc * 1.125) {
210 /* Moderate upsize; overallocate similar to list_resize() */
211 alloc = size + (size >> 3) + (size < 9 ? 3 : 6);
212 }
213 else {
214 /* Major upsize; resize up to exact size */
215 alloc = size + 1;
216 }
217
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000218 sval = PyMem_Realloc(((PyByteArrayObject *)self)->ob_bytes, alloc);
219 if (sval == NULL) {
220 PyErr_NoMemory();
221 return -1;
222 }
223
224 ((PyByteArrayObject *)self)->ob_bytes = sval;
225 Py_SIZE(self) = size;
226 ((PyByteArrayObject *)self)->ob_alloc = alloc;
227 ((PyByteArrayObject *)self)->ob_bytes[size] = '\0'; /* Trailing null byte */
228
229 return 0;
230}
231
232PyObject *
233PyByteArray_Concat(PyObject *a, PyObject *b)
234{
235 Py_ssize_t size;
236 Py_buffer va, vb;
237 PyByteArrayObject *result = NULL;
238
239 va.len = -1;
240 vb.len = -1;
241 if (_getbuffer(a, &va) < 0 ||
242 _getbuffer(b, &vb) < 0) {
243 PyErr_Format(PyExc_TypeError, "can't concat %.100s to %.100s",
244 Py_TYPE(a)->tp_name, Py_TYPE(b)->tp_name);
245 goto done;
246 }
247
248 size = va.len + vb.len;
249 if (size < 0) {
Benjamin Petersone0124bd2009-03-09 21:04:33 +0000250 PyErr_NoMemory();
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000251 goto done;
252 }
253
254 result = (PyByteArrayObject *) PyByteArray_FromStringAndSize(NULL, size);
255 if (result != NULL) {
256 memcpy(result->ob_bytes, va.buf, va.len);
257 memcpy(result->ob_bytes + va.len, vb.buf, vb.len);
258 }
259
260 done:
261 if (va.len != -1)
Martin v. Löwis423be952008-08-13 15:53:07 +0000262 PyBuffer_Release(&va);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000263 if (vb.len != -1)
Martin v. Löwis423be952008-08-13 15:53:07 +0000264 PyBuffer_Release(&vb);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000265 return (PyObject *)result;
266}
267
268/* Functions stuffed into the type object */
269
270static Py_ssize_t
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000271bytearray_length(PyByteArrayObject *self)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000272{
273 return Py_SIZE(self);
274}
275
276static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000277bytearray_iconcat(PyByteArrayObject *self, PyObject *other)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000278{
279 Py_ssize_t mysize;
280 Py_ssize_t size;
281 Py_buffer vo;
282
283 if (_getbuffer(other, &vo) < 0) {
284 PyErr_Format(PyExc_TypeError, "can't concat %.100s to %.100s",
285 Py_TYPE(other)->tp_name, Py_TYPE(self)->tp_name);
286 return NULL;
287 }
288
289 mysize = Py_SIZE(self);
290 size = mysize + vo.len;
291 if (size < 0) {
Martin v. Löwis423be952008-08-13 15:53:07 +0000292 PyBuffer_Release(&vo);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000293 return PyErr_NoMemory();
294 }
295 if (size < self->ob_alloc) {
296 Py_SIZE(self) = size;
297 self->ob_bytes[Py_SIZE(self)] = '\0'; /* Trailing null byte */
298 }
299 else if (PyByteArray_Resize((PyObject *)self, size) < 0) {
Martin v. Löwis423be952008-08-13 15:53:07 +0000300 PyBuffer_Release(&vo);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000301 return NULL;
302 }
303 memcpy(self->ob_bytes + mysize, vo.buf, vo.len);
Martin v. Löwis423be952008-08-13 15:53:07 +0000304 PyBuffer_Release(&vo);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000305 Py_INCREF(self);
306 return (PyObject *)self;
307}
308
309static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000310bytearray_repeat(PyByteArrayObject *self, Py_ssize_t count)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000311{
312 PyByteArrayObject *result;
313 Py_ssize_t mysize;
314 Py_ssize_t size;
315
316 if (count < 0)
317 count = 0;
318 mysize = Py_SIZE(self);
319 size = mysize * count;
320 if (count != 0 && size / count != mysize)
321 return PyErr_NoMemory();
322 result = (PyByteArrayObject *)PyByteArray_FromStringAndSize(NULL, size);
323 if (result != NULL && size != 0) {
324 if (mysize == 1)
325 memset(result->ob_bytes, self->ob_bytes[0], size);
326 else {
327 Py_ssize_t i;
328 for (i = 0; i < count; i++)
329 memcpy(result->ob_bytes + i*mysize, self->ob_bytes, mysize);
330 }
331 }
332 return (PyObject *)result;
333}
334
335static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000336bytearray_irepeat(PyByteArrayObject *self, Py_ssize_t count)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000337{
338 Py_ssize_t mysize;
339 Py_ssize_t size;
340
341 if (count < 0)
342 count = 0;
343 mysize = Py_SIZE(self);
344 size = mysize * count;
345 if (count != 0 && size / count != mysize)
346 return PyErr_NoMemory();
347 if (size < self->ob_alloc) {
348 Py_SIZE(self) = size;
349 self->ob_bytes[Py_SIZE(self)] = '\0'; /* Trailing null byte */
350 }
351 else if (PyByteArray_Resize((PyObject *)self, size) < 0)
352 return NULL;
353
354 if (mysize == 1)
355 memset(self->ob_bytes, self->ob_bytes[0], size);
356 else {
357 Py_ssize_t i;
358 for (i = 1; i < count; i++)
359 memcpy(self->ob_bytes + i*mysize, self->ob_bytes, mysize);
360 }
361
362 Py_INCREF(self);
363 return (PyObject *)self;
364}
365
366static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000367bytearray_getitem(PyByteArrayObject *self, Py_ssize_t i)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000368{
369 if (i < 0)
370 i += Py_SIZE(self);
371 if (i < 0 || i >= Py_SIZE(self)) {
372 PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
373 return NULL;
374 }
375 return PyLong_FromLong((unsigned char)(self->ob_bytes[i]));
376}
377
378static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000379bytearray_subscript(PyByteArrayObject *self, PyObject *index)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000380{
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000381 if (PyIndex_Check(index)) {
382 Py_ssize_t i = PyNumber_AsSsize_t(index, PyExc_IndexError);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000383
384 if (i == -1 && PyErr_Occurred())
385 return NULL;
386
387 if (i < 0)
388 i += PyByteArray_GET_SIZE(self);
389
390 if (i < 0 || i >= Py_SIZE(self)) {
391 PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
392 return NULL;
393 }
394 return PyLong_FromLong((unsigned char)(self->ob_bytes[i]));
395 }
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000396 else if (PySlice_Check(index)) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000397 Py_ssize_t start, stop, step, slicelength, cur, i;
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000398 if (PySlice_GetIndicesEx((PySliceObject *)index,
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000399 PyByteArray_GET_SIZE(self),
400 &start, &stop, &step, &slicelength) < 0) {
401 return NULL;
402 }
403
404 if (slicelength <= 0)
405 return PyByteArray_FromStringAndSize("", 0);
406 else if (step == 1) {
407 return PyByteArray_FromStringAndSize(self->ob_bytes + start,
408 slicelength);
409 }
410 else {
411 char *source_buf = PyByteArray_AS_STRING(self);
Alexandre Vassalottie2641f42009-04-03 06:38:02 +0000412 char *result_buf;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000413 PyObject *result;
414
Alexandre Vassalottie2641f42009-04-03 06:38:02 +0000415 result = PyByteArray_FromStringAndSize(NULL, slicelength);
416 if (result == NULL)
417 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000418
Alexandre Vassalottie2641f42009-04-03 06:38:02 +0000419 result_buf = PyByteArray_AS_STRING(result);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000420 for (cur = start, i = 0; i < slicelength;
421 cur += step, i++) {
422 result_buf[i] = source_buf[cur];
423 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000424 return result;
425 }
426 }
427 else {
428 PyErr_SetString(PyExc_TypeError, "bytearray indices must be integers");
429 return NULL;
430 }
431}
432
433static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000434bytearray_setslice(PyByteArrayObject *self, Py_ssize_t lo, Py_ssize_t hi,
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000435 PyObject *values)
436{
437 Py_ssize_t avail, needed;
438 void *bytes;
439 Py_buffer vbytes;
440 int res = 0;
441
442 vbytes.len = -1;
443 if (values == (PyObject *)self) {
444 /* Make a copy and call this function recursively */
445 int err;
446 values = PyByteArray_FromObject(values);
447 if (values == NULL)
448 return -1;
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000449 err = bytearray_setslice(self, lo, hi, values);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000450 Py_DECREF(values);
451 return err;
452 }
453 if (values == NULL) {
454 /* del b[lo:hi] */
455 bytes = NULL;
456 needed = 0;
457 }
458 else {
459 if (_getbuffer(values, &vbytes) < 0) {
460 PyErr_Format(PyExc_TypeError,
Georg Brandl3dbca812008-07-23 16:10:53 +0000461 "can't set bytearray slice from %.100s",
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000462 Py_TYPE(values)->tp_name);
463 return -1;
464 }
465 needed = vbytes.len;
466 bytes = vbytes.buf;
467 }
468
469 if (lo < 0)
470 lo = 0;
471 if (hi < lo)
472 hi = lo;
473 if (hi > Py_SIZE(self))
474 hi = Py_SIZE(self);
475
476 avail = hi - lo;
477 if (avail < 0)
478 lo = hi = avail = 0;
479
480 if (avail != needed) {
481 if (avail > needed) {
Antoine Pitrou5504e892008-12-06 21:27:53 +0000482 if (!_canresize(self)) {
483 res = -1;
484 goto finish;
485 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000486 /*
487 0 lo hi old_size
488 | |<----avail----->|<-----tomove------>|
489 | |<-needed->|<-----tomove------>|
490 0 lo new_hi new_size
491 */
492 memmove(self->ob_bytes + lo + needed, self->ob_bytes + hi,
493 Py_SIZE(self) - hi);
494 }
495 /* XXX(nnorwitz): need to verify this can't overflow! */
496 if (PyByteArray_Resize((PyObject *)self,
497 Py_SIZE(self) + needed - avail) < 0) {
498 res = -1;
499 goto finish;
500 }
501 if (avail < needed) {
502 /*
503 0 lo hi old_size
504 | |<-avail->|<-----tomove------>|
505 | |<----needed---->|<-----tomove------>|
506 0 lo new_hi new_size
507 */
508 memmove(self->ob_bytes + lo + needed, self->ob_bytes + hi,
509 Py_SIZE(self) - lo - needed);
510 }
511 }
512
513 if (needed > 0)
514 memcpy(self->ob_bytes + lo, bytes, needed);
515
516
517 finish:
518 if (vbytes.len != -1)
Martin v. Löwis423be952008-08-13 15:53:07 +0000519 PyBuffer_Release(&vbytes);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000520 return res;
521}
522
523static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000524bytearray_setitem(PyByteArrayObject *self, Py_ssize_t i, PyObject *value)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000525{
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000526 int ival;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000527
528 if (i < 0)
529 i += Py_SIZE(self);
530
531 if (i < 0 || i >= Py_SIZE(self)) {
532 PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
533 return -1;
534 }
535
536 if (value == NULL)
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000537 return bytearray_setslice(self, i, i+1, NULL);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000538
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000539 if (!_getbytevalue(value, &ival))
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000540 return -1;
541
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000542 self->ob_bytes[i] = ival;
543 return 0;
544}
545
546static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000547bytearray_ass_subscript(PyByteArrayObject *self, PyObject *index, PyObject *values)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000548{
549 Py_ssize_t start, stop, step, slicelen, needed;
550 char *bytes;
551
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000552 if (PyIndex_Check(index)) {
553 Py_ssize_t i = PyNumber_AsSsize_t(index, PyExc_IndexError);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000554
555 if (i == -1 && PyErr_Occurred())
556 return -1;
557
558 if (i < 0)
559 i += PyByteArray_GET_SIZE(self);
560
561 if (i < 0 || i >= Py_SIZE(self)) {
562 PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
563 return -1;
564 }
565
566 if (values == NULL) {
567 /* Fall through to slice assignment */
568 start = i;
569 stop = i + 1;
570 step = 1;
571 slicelen = 1;
572 }
573 else {
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000574 int ival;
575 if (!_getbytevalue(values, &ival))
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000576 return -1;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000577 self->ob_bytes[i] = (char)ival;
578 return 0;
579 }
580 }
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000581 else if (PySlice_Check(index)) {
582 if (PySlice_GetIndicesEx((PySliceObject *)index,
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000583 PyByteArray_GET_SIZE(self),
584 &start, &stop, &step, &slicelen) < 0) {
585 return -1;
586 }
587 }
588 else {
589 PyErr_SetString(PyExc_TypeError, "bytearray indices must be integer");
590 return -1;
591 }
592
593 if (values == NULL) {
594 bytes = NULL;
595 needed = 0;
596 }
597 else if (values == (PyObject *)self || !PyByteArray_Check(values)) {
Georg Brandlf6c8fd62011-02-25 09:48:21 +0000598 /* Make a copy and call this function recursively */
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000599 int err;
600 values = PyByteArray_FromObject(values);
601 if (values == NULL)
602 return -1;
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000603 err = bytearray_ass_subscript(self, index, values);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000604 Py_DECREF(values);
605 return err;
606 }
607 else {
608 assert(PyByteArray_Check(values));
609 bytes = ((PyByteArrayObject *)values)->ob_bytes;
610 needed = Py_SIZE(values);
611 }
612 /* Make sure b[5:2] = ... inserts before 5, not before 2. */
613 if ((step < 0 && start < stop) ||
614 (step > 0 && start > stop))
615 stop = start;
616 if (step == 1) {
617 if (slicelen != needed) {
Antoine Pitrou5504e892008-12-06 21:27:53 +0000618 if (!_canresize(self))
619 return -1;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000620 if (slicelen > needed) {
621 /*
622 0 start stop old_size
623 | |<---slicelen--->|<-----tomove------>|
624 | |<-needed->|<-----tomove------>|
625 0 lo new_hi new_size
626 */
627 memmove(self->ob_bytes + start + needed, self->ob_bytes + stop,
628 Py_SIZE(self) - stop);
629 }
630 if (PyByteArray_Resize((PyObject *)self,
631 Py_SIZE(self) + needed - slicelen) < 0)
632 return -1;
633 if (slicelen < needed) {
634 /*
635 0 lo hi old_size
636 | |<-avail->|<-----tomove------>|
637 | |<----needed---->|<-----tomove------>|
638 0 lo new_hi new_size
639 */
640 memmove(self->ob_bytes + start + needed, self->ob_bytes + stop,
641 Py_SIZE(self) - start - needed);
642 }
643 }
644
645 if (needed > 0)
646 memcpy(self->ob_bytes + start, bytes, needed);
647
648 return 0;
649 }
650 else {
651 if (needed == 0) {
652 /* Delete slice */
Mark Dickinsona53f2c92010-01-29 17:29:21 +0000653 size_t cur;
654 Py_ssize_t i;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000655
Antoine Pitrou5504e892008-12-06 21:27:53 +0000656 if (!_canresize(self))
657 return -1;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000658 if (step < 0) {
659 stop = start + 1;
660 start = stop + step * (slicelen - 1) - 1;
661 step = -step;
662 }
663 for (cur = start, i = 0;
664 i < slicelen; cur += step, i++) {
665 Py_ssize_t lim = step - 1;
666
Mark Dickinson5a41a4c2010-02-14 13:09:30 +0000667 if (cur + step >= (size_t)PyByteArray_GET_SIZE(self))
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000668 lim = PyByteArray_GET_SIZE(self) - cur - 1;
669
670 memmove(self->ob_bytes + cur - i,
671 self->ob_bytes + cur + 1, lim);
672 }
673 /* Move the tail of the bytes, in one chunk */
674 cur = start + slicelen*step;
Mark Dickinson5a41a4c2010-02-14 13:09:30 +0000675 if (cur < (size_t)PyByteArray_GET_SIZE(self)) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000676 memmove(self->ob_bytes + cur - slicelen,
677 self->ob_bytes + cur,
678 PyByteArray_GET_SIZE(self) - cur);
679 }
680 if (PyByteArray_Resize((PyObject *)self,
681 PyByteArray_GET_SIZE(self) - slicelen) < 0)
682 return -1;
683
684 return 0;
685 }
686 else {
687 /* Assign slice */
688 Py_ssize_t cur, i;
689
690 if (needed != slicelen) {
691 PyErr_Format(PyExc_ValueError,
692 "attempt to assign bytes of size %zd "
693 "to extended slice of size %zd",
694 needed, slicelen);
695 return -1;
696 }
697 for (cur = start, i = 0; i < slicelen; cur += step, i++)
698 self->ob_bytes[cur] = bytes[i];
699 return 0;
700 }
701 }
702}
703
704static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000705bytearray_init(PyByteArrayObject *self, PyObject *args, PyObject *kwds)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000706{
707 static char *kwlist[] = {"source", "encoding", "errors", 0};
708 PyObject *arg = NULL;
709 const char *encoding = NULL;
710 const char *errors = NULL;
711 Py_ssize_t count;
712 PyObject *it;
713 PyObject *(*iternext)(PyObject *);
714
715 if (Py_SIZE(self) != 0) {
716 /* Empty previous contents (yes, do this first of all!) */
717 if (PyByteArray_Resize((PyObject *)self, 0) < 0)
718 return -1;
719 }
720
721 /* Parse arguments */
Georg Brandl3dbca812008-07-23 16:10:53 +0000722 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|Oss:bytearray", kwlist,
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000723 &arg, &encoding, &errors))
724 return -1;
725
726 /* Make a quick exit if no first argument */
727 if (arg == NULL) {
728 if (encoding != NULL || errors != NULL) {
729 PyErr_SetString(PyExc_TypeError,
730 "encoding or errors without sequence argument");
731 return -1;
732 }
733 return 0;
734 }
735
736 if (PyUnicode_Check(arg)) {
737 /* Encode via the codec registry */
738 PyObject *encoded, *new;
739 if (encoding == NULL) {
740 PyErr_SetString(PyExc_TypeError,
741 "string argument without an encoding");
742 return -1;
743 }
Marc-André Lemburgb2750b52008-06-06 12:18:17 +0000744 encoded = PyUnicode_AsEncodedString(arg, encoding, errors);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000745 if (encoded == NULL)
746 return -1;
747 assert(PyBytes_Check(encoded));
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000748 new = bytearray_iconcat(self, encoded);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000749 Py_DECREF(encoded);
750 if (new == NULL)
751 return -1;
752 Py_DECREF(new);
753 return 0;
754 }
755
756 /* If it's not unicode, there can't be encoding or errors */
757 if (encoding != NULL || errors != NULL) {
758 PyErr_SetString(PyExc_TypeError,
759 "encoding or errors without a string argument");
760 return -1;
761 }
762
763 /* Is it an int? */
Benjamin Peterson0ff8a502010-04-16 23:19:11 +0000764 count = PyNumber_AsSsize_t(arg, PyExc_OverflowError);
765 if (count == -1 && PyErr_Occurred()) {
766 if (PyErr_ExceptionMatches(PyExc_OverflowError))
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000767 return -1;
Benjamin Peterson0ff8a502010-04-16 23:19:11 +0000768 PyErr_Clear();
769 }
770 else if (count < 0) {
771 PyErr_SetString(PyExc_ValueError, "negative count");
772 return -1;
773 }
774 else {
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000775 if (count > 0) {
776 if (PyByteArray_Resize((PyObject *)self, count))
777 return -1;
778 memset(self->ob_bytes, 0, count);
779 }
780 return 0;
781 }
782
783 /* Use the buffer API */
784 if (PyObject_CheckBuffer(arg)) {
785 Py_ssize_t size;
786 Py_buffer view;
787 if (PyObject_GetBuffer(arg, &view, PyBUF_FULL_RO) < 0)
788 return -1;
789 size = view.len;
790 if (PyByteArray_Resize((PyObject *)self, size) < 0) goto fail;
791 if (PyBuffer_ToContiguous(self->ob_bytes, &view, size, 'C') < 0)
792 goto fail;
Martin v. Löwis423be952008-08-13 15:53:07 +0000793 PyBuffer_Release(&view);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000794 return 0;
795 fail:
Martin v. Löwis423be952008-08-13 15:53:07 +0000796 PyBuffer_Release(&view);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000797 return -1;
798 }
799
800 /* XXX Optimize this if the arguments is a list, tuple */
801
802 /* Get the iterator */
803 it = PyObject_GetIter(arg);
804 if (it == NULL)
805 return -1;
806 iternext = *Py_TYPE(it)->tp_iternext;
807
808 /* Run the iterator to exhaustion */
809 for (;;) {
810 PyObject *item;
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000811 int rc, value;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000812
813 /* Get the next item */
814 item = iternext(it);
815 if (item == NULL) {
816 if (PyErr_Occurred()) {
817 if (!PyErr_ExceptionMatches(PyExc_StopIteration))
818 goto error;
819 PyErr_Clear();
820 }
821 break;
822 }
823
824 /* Interpret it as an int (__index__) */
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000825 rc = _getbytevalue(item, &value);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000826 Py_DECREF(item);
Georg Brandl9a54d7c2008-07-16 23:15:30 +0000827 if (!rc)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000828 goto error;
829
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000830 /* Append the byte */
831 if (Py_SIZE(self) < self->ob_alloc)
832 Py_SIZE(self)++;
833 else if (PyByteArray_Resize((PyObject *)self, Py_SIZE(self)+1) < 0)
834 goto error;
835 self->ob_bytes[Py_SIZE(self)-1] = value;
836 }
837
838 /* Clean up and return success */
839 Py_DECREF(it);
840 return 0;
841
842 error:
843 /* Error handling when it != NULL */
844 Py_DECREF(it);
845 return -1;
846}
847
848/* Mostly copied from string_repr, but without the
849 "smart quote" functionality. */
850static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000851bytearray_repr(PyByteArrayObject *self)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000852{
853 static const char *hexdigits = "0123456789abcdef";
854 const char *quote_prefix = "bytearray(b";
855 const char *quote_postfix = ")";
856 Py_ssize_t length = Py_SIZE(self);
857 /* 14 == strlen(quote_prefix) + 2 + strlen(quote_postfix) */
Mark Dickinson5a41a4c2010-02-14 13:09:30 +0000858 size_t newsize;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000859 PyObject *v;
Mark Dickinson5a41a4c2010-02-14 13:09:30 +0000860 if (length > (PY_SSIZE_T_MAX - 14) / 4) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000861 PyErr_SetString(PyExc_OverflowError,
862 "bytearray object is too large to make repr");
863 return NULL;
864 }
Mark Dickinson5a41a4c2010-02-14 13:09:30 +0000865 newsize = 14 + 4 * length;
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000866 v = PyUnicode_FromUnicode(NULL, newsize);
867 if (v == NULL) {
868 return NULL;
869 }
870 else {
871 register Py_ssize_t i;
872 register Py_UNICODE c;
873 register Py_UNICODE *p;
874 int quote;
875
876 /* Figure out which quote to use; single is preferred */
877 quote = '\'';
878 {
879 char *test, *start;
880 start = PyByteArray_AS_STRING(self);
881 for (test = start; test < start+length; ++test) {
882 if (*test == '"') {
883 quote = '\''; /* back to single */
884 goto decided;
885 }
886 else if (*test == '\'')
887 quote = '"';
888 }
889 decided:
890 ;
891 }
892
893 p = PyUnicode_AS_UNICODE(v);
894 while (*quote_prefix)
895 *p++ = *quote_prefix++;
896 *p++ = quote;
897
898 for (i = 0; i < length; i++) {
899 /* There's at least enough room for a hex escape
900 and a closing quote. */
901 assert(newsize - (p - PyUnicode_AS_UNICODE(v)) >= 5);
902 c = self->ob_bytes[i];
903 if (c == '\'' || c == '\\')
904 *p++ = '\\', *p++ = c;
905 else if (c == '\t')
906 *p++ = '\\', *p++ = 't';
907 else if (c == '\n')
908 *p++ = '\\', *p++ = 'n';
909 else if (c == '\r')
910 *p++ = '\\', *p++ = 'r';
911 else if (c == 0)
912 *p++ = '\\', *p++ = 'x', *p++ = '0', *p++ = '0';
913 else if (c < ' ' || c >= 0x7f) {
914 *p++ = '\\';
915 *p++ = 'x';
916 *p++ = hexdigits[(c & 0xf0) >> 4];
917 *p++ = hexdigits[c & 0xf];
918 }
919 else
920 *p++ = c;
921 }
922 assert(newsize - (p - PyUnicode_AS_UNICODE(v)) >= 1);
923 *p++ = quote;
924 while (*quote_postfix) {
925 *p++ = *quote_postfix++;
926 }
927 *p = '\0';
928 if (PyUnicode_Resize(&v, (p - PyUnicode_AS_UNICODE(v)))) {
929 Py_DECREF(v);
930 return NULL;
931 }
932 return v;
933 }
934}
935
936static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000937bytearray_str(PyObject *op)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000938{
Ezio Melotti42da6632011-03-15 05:18:48 +0200939 if (Py_BytesWarningFlag) {
940 if (PyErr_WarnEx(PyExc_BytesWarning,
941 "str() on a bytearray instance", 1))
942 return NULL;
943 }
944 return bytearray_repr((PyByteArrayObject*)op);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000945}
946
947static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000948bytearray_richcompare(PyObject *self, PyObject *other, int op)
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000949{
950 Py_ssize_t self_size, other_size;
951 Py_buffer self_bytes, other_bytes;
952 PyObject *res;
953 Py_ssize_t minsize;
954 int cmp;
955
956 /* Bytes can be compared to anything that supports the (binary)
957 buffer API. Except that a comparison with Unicode is always an
958 error, even if the comparison is for equality. */
959 if (PyObject_IsInstance(self, (PyObject*)&PyUnicode_Type) ||
960 PyObject_IsInstance(other, (PyObject*)&PyUnicode_Type)) {
Barry Warsaw9e9dcd62008-10-17 01:50:37 +0000961 if (Py_BytesWarningFlag && (op == Py_EQ || op == Py_NE)) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000962 if (PyErr_WarnEx(PyExc_BytesWarning,
Georg Brandle5d68ac2008-06-04 11:30:26 +0000963 "Comparison between bytearray and string", 1))
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000964 return NULL;
965 }
966
967 Py_INCREF(Py_NotImplemented);
968 return Py_NotImplemented;
969 }
970
971 self_size = _getbuffer(self, &self_bytes);
972 if (self_size < 0) {
973 PyErr_Clear();
974 Py_INCREF(Py_NotImplemented);
975 return Py_NotImplemented;
976 }
977
978 other_size = _getbuffer(other, &other_bytes);
979 if (other_size < 0) {
980 PyErr_Clear();
Martin v. Löwis423be952008-08-13 15:53:07 +0000981 PyBuffer_Release(&self_bytes);
Christian Heimes2c9c7a52008-05-26 13:42:13 +0000982 Py_INCREF(Py_NotImplemented);
983 return Py_NotImplemented;
984 }
985
986 if (self_size != other_size && (op == Py_EQ || op == Py_NE)) {
987 /* Shortcut: if the lengths differ, the objects differ */
988 cmp = (op == Py_NE);
989 }
990 else {
991 minsize = self_size;
992 if (other_size < minsize)
993 minsize = other_size;
994
995 cmp = memcmp(self_bytes.buf, other_bytes.buf, minsize);
996 /* In ISO C, memcmp() guarantees to use unsigned bytes! */
997
998 if (cmp == 0) {
999 if (self_size < other_size)
1000 cmp = -1;
1001 else if (self_size > other_size)
1002 cmp = 1;
1003 }
1004
1005 switch (op) {
1006 case Py_LT: cmp = cmp < 0; break;
1007 case Py_LE: cmp = cmp <= 0; break;
1008 case Py_EQ: cmp = cmp == 0; break;
1009 case Py_NE: cmp = cmp != 0; break;
1010 case Py_GT: cmp = cmp > 0; break;
1011 case Py_GE: cmp = cmp >= 0; break;
1012 }
1013 }
1014
1015 res = cmp ? Py_True : Py_False;
Martin v. Löwis423be952008-08-13 15:53:07 +00001016 PyBuffer_Release(&self_bytes);
1017 PyBuffer_Release(&other_bytes);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001018 Py_INCREF(res);
1019 return res;
1020}
1021
1022static void
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001023bytearray_dealloc(PyByteArrayObject *self)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001024{
Benjamin Petersone0124bd2009-03-09 21:04:33 +00001025 if (self->ob_exports > 0) {
1026 PyErr_SetString(PyExc_SystemError,
Benjamin Petersonef3e4c22009-04-11 19:48:14 +00001027 "deallocated bytearray object has exported buffers");
Benjamin Petersone0124bd2009-03-09 21:04:33 +00001028 PyErr_Print();
1029 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001030 if (self->ob_bytes != 0) {
1031 PyMem_Free(self->ob_bytes);
1032 }
1033 Py_TYPE(self)->tp_free((PyObject *)self);
1034}
1035
1036
1037/* -------------------------------------------------------------------- */
1038/* Methods */
1039
1040#define STRINGLIB_CHAR char
1041#define STRINGLIB_CMP memcmp
1042#define STRINGLIB_LEN PyByteArray_GET_SIZE
1043#define STRINGLIB_STR PyByteArray_AS_STRING
1044#define STRINGLIB_NEW PyByteArray_FromStringAndSize
1045#define STRINGLIB_EMPTY nullbytes
1046#define STRINGLIB_CHECK_EXACT PyByteArray_CheckExact
1047#define STRINGLIB_MUTABLE 1
Benjamin Petersona786b022008-08-25 21:05:21 +00001048#define FROM_BYTEARRAY 1
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001049
1050#include "stringlib/fastsearch.h"
1051#include "stringlib/count.h"
1052#include "stringlib/find.h"
1053#include "stringlib/partition.h"
1054#include "stringlib/ctype.h"
1055#include "stringlib/transmogrify.h"
1056
1057
1058/* The following Py_LOCAL_INLINE and Py_LOCAL functions
1059were copied from the old char* style string object. */
1060
1061Py_LOCAL_INLINE(void)
1062_adjust_indices(Py_ssize_t *start, Py_ssize_t *end, Py_ssize_t len)
1063{
1064 if (*end > len)
1065 *end = len;
1066 else if (*end < 0)
1067 *end += len;
1068 if (*end < 0)
1069 *end = 0;
1070 if (*start < 0)
1071 *start += len;
1072 if (*start < 0)
1073 *start = 0;
1074}
1075
1076
1077Py_LOCAL_INLINE(Py_ssize_t)
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001078bytearray_find_internal(PyByteArrayObject *self, PyObject *args, int dir)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001079{
1080 PyObject *subobj;
1081 Py_buffer subbuf;
1082 Py_ssize_t start=0, end=PY_SSIZE_T_MAX;
1083 Py_ssize_t res;
1084
Jesus Ceaac451502011-04-20 17:09:23 +02001085 if (!stringlib_parse_args_finds("find/rfind/index/rindex",
1086 args, &subobj, &start, &end))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001087 return -2;
1088 if (_getbuffer(subobj, &subbuf) < 0)
1089 return -2;
1090 if (dir > 0)
1091 res = stringlib_find_slice(
1092 PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
1093 subbuf.buf, subbuf.len, start, end);
1094 else
1095 res = stringlib_rfind_slice(
1096 PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
1097 subbuf.buf, subbuf.len, start, end);
Martin v. Löwis423be952008-08-13 15:53:07 +00001098 PyBuffer_Release(&subbuf);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001099 return res;
1100}
1101
1102PyDoc_STRVAR(find__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001103"B.find(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001104\n\
1105Return the lowest index in B where subsection sub is found,\n\
1106such that sub is contained within s[start,end]. Optional\n\
1107arguments start and end are interpreted as in slice notation.\n\
1108\n\
1109Return -1 on failure.");
1110
1111static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001112bytearray_find(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001113{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001114 Py_ssize_t result = bytearray_find_internal(self, args, +1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001115 if (result == -2)
1116 return NULL;
1117 return PyLong_FromSsize_t(result);
1118}
1119
1120PyDoc_STRVAR(count__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001121"B.count(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001122\n\
1123Return the number of non-overlapping occurrences of subsection sub in\n\
1124bytes B[start:end]. Optional arguments start and end are interpreted\n\
1125as in slice notation.");
1126
1127static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001128bytearray_count(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001129{
1130 PyObject *sub_obj;
1131 const char *str = PyByteArray_AS_STRING(self);
1132 Py_ssize_t start = 0, end = PY_SSIZE_T_MAX;
1133 Py_buffer vsub;
1134 PyObject *count_obj;
1135
Jesus Ceaac451502011-04-20 17:09:23 +02001136 if (!stringlib_parse_args_finds("count", args, &sub_obj, &start, &end))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001137 return NULL;
1138
1139 if (_getbuffer(sub_obj, &vsub) < 0)
1140 return NULL;
1141
1142 _adjust_indices(&start, &end, PyByteArray_GET_SIZE(self));
1143
1144 count_obj = PyLong_FromSsize_t(
1145 stringlib_count(str + start, end - start, vsub.buf, vsub.len)
1146 );
Martin v. Löwis423be952008-08-13 15:53:07 +00001147 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001148 return count_obj;
1149}
1150
1151
1152PyDoc_STRVAR(index__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001153"B.index(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001154\n\
1155Like B.find() but raise ValueError when the subsection is not found.");
1156
1157static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001158bytearray_index(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001159{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001160 Py_ssize_t result = bytearray_find_internal(self, args, +1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001161 if (result == -2)
1162 return NULL;
1163 if (result == -1) {
1164 PyErr_SetString(PyExc_ValueError,
1165 "subsection not found");
1166 return NULL;
1167 }
1168 return PyLong_FromSsize_t(result);
1169}
1170
1171
1172PyDoc_STRVAR(rfind__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001173"B.rfind(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001174\n\
1175Return the highest index in B where subsection sub is found,\n\
1176such that sub is contained within s[start,end]. Optional\n\
1177arguments start and end are interpreted as in slice notation.\n\
1178\n\
1179Return -1 on failure.");
1180
1181static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001182bytearray_rfind(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001183{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001184 Py_ssize_t result = bytearray_find_internal(self, args, -1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001185 if (result == -2)
1186 return NULL;
1187 return PyLong_FromSsize_t(result);
1188}
1189
1190
1191PyDoc_STRVAR(rindex__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001192"B.rindex(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001193\n\
1194Like B.rfind() but raise ValueError when the subsection is not found.");
1195
1196static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001197bytearray_rindex(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001198{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001199 Py_ssize_t result = bytearray_find_internal(self, args, -1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001200 if (result == -2)
1201 return NULL;
1202 if (result == -1) {
1203 PyErr_SetString(PyExc_ValueError,
1204 "subsection not found");
1205 return NULL;
1206 }
1207 return PyLong_FromSsize_t(result);
1208}
1209
1210
1211static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001212bytearray_contains(PyObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001213{
1214 Py_ssize_t ival = PyNumber_AsSsize_t(arg, PyExc_ValueError);
1215 if (ival == -1 && PyErr_Occurred()) {
1216 Py_buffer varg;
Antoine Pitroubc760d92010-08-15 17:46:50 +00001217 Py_ssize_t pos;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001218 PyErr_Clear();
1219 if (_getbuffer(arg, &varg) < 0)
1220 return -1;
1221 pos = stringlib_find(PyByteArray_AS_STRING(self), Py_SIZE(self),
1222 varg.buf, varg.len, 0);
Martin v. Löwis423be952008-08-13 15:53:07 +00001223 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001224 return pos >= 0;
1225 }
1226 if (ival < 0 || ival >= 256) {
1227 PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
1228 return -1;
1229 }
1230
Antoine Pitroubc760d92010-08-15 17:46:50 +00001231 return memchr(PyByteArray_AS_STRING(self), (int) ival, Py_SIZE(self)) != NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001232}
1233
1234
1235/* Matches the end (direction >= 0) or start (direction < 0) of self
1236 * against substr, using the start and end arguments. Returns
1237 * -1 on error, 0 if not found and 1 if found.
1238 */
1239Py_LOCAL(int)
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001240_bytearray_tailmatch(PyByteArrayObject *self, PyObject *substr, Py_ssize_t start,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001241 Py_ssize_t end, int direction)
1242{
1243 Py_ssize_t len = PyByteArray_GET_SIZE(self);
1244 const char* str;
1245 Py_buffer vsubstr;
1246 int rv = 0;
1247
1248 str = PyByteArray_AS_STRING(self);
1249
1250 if (_getbuffer(substr, &vsubstr) < 0)
1251 return -1;
1252
1253 _adjust_indices(&start, &end, len);
1254
1255 if (direction < 0) {
1256 /* startswith */
1257 if (start+vsubstr.len > len) {
1258 goto done;
1259 }
1260 } else {
1261 /* endswith */
1262 if (end-start < vsubstr.len || start > len) {
1263 goto done;
1264 }
1265
1266 if (end-vsubstr.len > start)
1267 start = end - vsubstr.len;
1268 }
1269 if (end-start >= vsubstr.len)
1270 rv = ! memcmp(str+start, vsubstr.buf, vsubstr.len);
1271
1272done:
Martin v. Löwis423be952008-08-13 15:53:07 +00001273 PyBuffer_Release(&vsubstr);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001274 return rv;
1275}
1276
1277
1278PyDoc_STRVAR(startswith__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001279"B.startswith(prefix[, start[, end]]) -> bool\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001280\n\
1281Return True if B starts with the specified prefix, False otherwise.\n\
1282With optional start, test B beginning at that position.\n\
1283With optional end, stop comparing B at that position.\n\
Ezio Melottiba42fd52011-04-26 06:09:45 +03001284prefix can also be a tuple of bytes to try.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001285
1286static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001287bytearray_startswith(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001288{
1289 Py_ssize_t start = 0;
1290 Py_ssize_t end = PY_SSIZE_T_MAX;
1291 PyObject *subobj;
1292 int result;
1293
Jesus Ceaac451502011-04-20 17:09:23 +02001294 if (!stringlib_parse_args_finds("startswith", args, &subobj, &start, &end))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001295 return NULL;
1296 if (PyTuple_Check(subobj)) {
1297 Py_ssize_t i;
1298 for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001299 result = _bytearray_tailmatch(self,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001300 PyTuple_GET_ITEM(subobj, i),
1301 start, end, -1);
1302 if (result == -1)
1303 return NULL;
1304 else if (result) {
1305 Py_RETURN_TRUE;
1306 }
1307 }
1308 Py_RETURN_FALSE;
1309 }
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001310 result = _bytearray_tailmatch(self, subobj, start, end, -1);
Ezio Melottiba42fd52011-04-26 06:09:45 +03001311 if (result == -1) {
1312 if (PyErr_ExceptionMatches(PyExc_TypeError))
1313 PyErr_Format(PyExc_TypeError, "startswith first arg must be bytes "
1314 "or a tuple of bytes, not %s", Py_TYPE(subobj)->tp_name);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001315 return NULL;
Ezio Melottiba42fd52011-04-26 06:09:45 +03001316 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001317 else
1318 return PyBool_FromLong(result);
1319}
1320
1321PyDoc_STRVAR(endswith__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001322"B.endswith(suffix[, start[, end]]) -> bool\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001323\n\
1324Return True if B ends with the specified suffix, False otherwise.\n\
1325With optional start, test B beginning at that position.\n\
1326With optional end, stop comparing B at that position.\n\
Ezio Melottiba42fd52011-04-26 06:09:45 +03001327suffix can also be a tuple of bytes to try.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001328
1329static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001330bytearray_endswith(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001331{
1332 Py_ssize_t start = 0;
1333 Py_ssize_t end = PY_SSIZE_T_MAX;
1334 PyObject *subobj;
1335 int result;
1336
Jesus Ceaac451502011-04-20 17:09:23 +02001337 if (!stringlib_parse_args_finds("endswith", args, &subobj, &start, &end))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001338 return NULL;
1339 if (PyTuple_Check(subobj)) {
1340 Py_ssize_t i;
1341 for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001342 result = _bytearray_tailmatch(self,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001343 PyTuple_GET_ITEM(subobj, i),
1344 start, end, +1);
1345 if (result == -1)
1346 return NULL;
1347 else if (result) {
1348 Py_RETURN_TRUE;
1349 }
1350 }
1351 Py_RETURN_FALSE;
1352 }
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001353 result = _bytearray_tailmatch(self, subobj, start, end, +1);
Ezio Melottiba42fd52011-04-26 06:09:45 +03001354 if (result == -1) {
1355 if (PyErr_ExceptionMatches(PyExc_TypeError))
1356 PyErr_Format(PyExc_TypeError, "endswith first arg must be bytes or "
1357 "a tuple of bytes, not %s", Py_TYPE(subobj)->tp_name);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001358 return NULL;
Ezio Melottiba42fd52011-04-26 06:09:45 +03001359 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001360 else
1361 return PyBool_FromLong(result);
1362}
1363
1364
1365PyDoc_STRVAR(translate__doc__,
1366"B.translate(table[, deletechars]) -> bytearray\n\
1367\n\
1368Return a copy of B, where all characters occurring in the\n\
1369optional argument deletechars are removed, and the remaining\n\
1370characters have been mapped through the given translation\n\
1371table, which must be a bytes object of length 256.");
1372
1373static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001374bytearray_translate(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001375{
1376 register char *input, *output;
1377 register const char *table;
Benjamin Peterson0f3641c2008-11-19 22:05:52 +00001378 register Py_ssize_t i, c;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001379 PyObject *input_obj = (PyObject*)self;
1380 const char *output_start;
1381 Py_ssize_t inlen;
Georg Brandlccc47b62008-12-28 11:44:14 +00001382 PyObject *result = NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001383 int trans_table[256];
Georg Brandlccc47b62008-12-28 11:44:14 +00001384 PyObject *tableobj = NULL, *delobj = NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001385 Py_buffer vtable, vdel;
1386
1387 if (!PyArg_UnpackTuple(args, "translate", 1, 2,
1388 &tableobj, &delobj))
1389 return NULL;
1390
Georg Brandlccc47b62008-12-28 11:44:14 +00001391 if (tableobj == Py_None) {
1392 table = NULL;
1393 tableobj = NULL;
1394 } else if (_getbuffer(tableobj, &vtable) < 0) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001395 return NULL;
Georg Brandlccc47b62008-12-28 11:44:14 +00001396 } else {
1397 if (vtable.len != 256) {
1398 PyErr_SetString(PyExc_ValueError,
1399 "translation table must be 256 characters long");
Georg Brandl069bcc32009-07-22 12:06:11 +00001400 PyBuffer_Release(&vtable);
1401 return NULL;
Georg Brandlccc47b62008-12-28 11:44:14 +00001402 }
1403 table = (const char*)vtable.buf;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001404 }
1405
1406 if (delobj != NULL) {
1407 if (_getbuffer(delobj, &vdel) < 0) {
Georg Brandl069bcc32009-07-22 12:06:11 +00001408 if (tableobj != NULL)
1409 PyBuffer_Release(&vtable);
1410 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001411 }
1412 }
1413 else {
1414 vdel.buf = NULL;
1415 vdel.len = 0;
1416 }
1417
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001418 inlen = PyByteArray_GET_SIZE(input_obj);
1419 result = PyByteArray_FromStringAndSize((char *)NULL, inlen);
1420 if (result == NULL)
1421 goto done;
1422 output_start = output = PyByteArray_AsString(result);
1423 input = PyByteArray_AS_STRING(input_obj);
1424
Georg Brandlccc47b62008-12-28 11:44:14 +00001425 if (vdel.len == 0 && table != NULL) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001426 /* If no deletions are required, use faster code */
1427 for (i = inlen; --i >= 0; ) {
1428 c = Py_CHARMASK(*input++);
Benjamin Peterson0f3641c2008-11-19 22:05:52 +00001429 *output++ = table[c];
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001430 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001431 goto done;
1432 }
Georg Brandlccc47b62008-12-28 11:44:14 +00001433
1434 if (table == NULL) {
1435 for (i = 0; i < 256; i++)
1436 trans_table[i] = Py_CHARMASK(i);
1437 } else {
1438 for (i = 0; i < 256; i++)
1439 trans_table[i] = Py_CHARMASK(table[i]);
1440 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001441
1442 for (i = 0; i < vdel.len; i++)
1443 trans_table[(int) Py_CHARMASK( ((unsigned char*)vdel.buf)[i] )] = -1;
1444
1445 for (i = inlen; --i >= 0; ) {
1446 c = Py_CHARMASK(*input++);
1447 if (trans_table[c] != -1)
1448 if (Py_CHARMASK(*output++ = (char)trans_table[c]) == c)
1449 continue;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001450 }
1451 /* Fix the size of the resulting string */
1452 if (inlen > 0)
1453 PyByteArray_Resize(result, output - output_start);
1454
1455done:
Georg Brandlccc47b62008-12-28 11:44:14 +00001456 if (tableobj != NULL)
1457 PyBuffer_Release(&vtable);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001458 if (delobj != NULL)
Martin v. Löwis423be952008-08-13 15:53:07 +00001459 PyBuffer_Release(&vdel);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001460 return result;
1461}
1462
1463
Georg Brandlabc38772009-04-12 15:51:51 +00001464static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001465bytearray_maketrans(PyObject *null, PyObject *args)
Georg Brandlabc38772009-04-12 15:51:51 +00001466{
Ezio Melotti42da6632011-03-15 05:18:48 +02001467 return _Py_bytes_maketrans(args);
Georg Brandlabc38772009-04-12 15:51:51 +00001468}
1469
1470
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001471#define FORWARD 1
1472#define REVERSE -1
1473
1474/* find and count characters and substrings */
1475
1476#define findchar(target, target_len, c) \
1477 ((char *)memchr((const void *)(target), c, target_len))
1478
1479/* Don't call if length < 2 */
1480#define Py_STRING_MATCH(target, offset, pattern, length) \
1481 (target[offset] == pattern[0] && \
1482 target[offset+length-1] == pattern[length-1] && \
1483 !memcmp(target+offset+1, pattern+1, length-2) )
1484
1485
Benjamin Peterson0f3641c2008-11-19 22:05:52 +00001486/* Bytes ops must return a string, create a copy */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001487Py_LOCAL(PyByteArrayObject *)
1488return_self(PyByteArrayObject *self)
1489{
Georg Brandl1e7217d2008-05-30 12:02:38 +00001490 /* always return a new bytearray */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001491 return (PyByteArrayObject *)PyByteArray_FromStringAndSize(
1492 PyByteArray_AS_STRING(self),
1493 PyByteArray_GET_SIZE(self));
1494}
1495
1496Py_LOCAL_INLINE(Py_ssize_t)
1497countchar(const char *target, Py_ssize_t target_len, char c, Py_ssize_t maxcount)
1498{
1499 Py_ssize_t count=0;
1500 const char *start=target;
1501 const char *end=target+target_len;
1502
1503 while ( (start=findchar(start, end-start, c)) != NULL ) {
1504 count++;
1505 if (count >= maxcount)
1506 break;
1507 start += 1;
1508 }
1509 return count;
1510}
1511
1512Py_LOCAL(Py_ssize_t)
1513findstring(const char *target, Py_ssize_t target_len,
1514 const char *pattern, Py_ssize_t pattern_len,
1515 Py_ssize_t start,
1516 Py_ssize_t end,
1517 int direction)
1518{
1519 if (start < 0) {
1520 start += target_len;
1521 if (start < 0)
1522 start = 0;
1523 }
1524 if (end > target_len) {
1525 end = target_len;
1526 } else if (end < 0) {
1527 end += target_len;
1528 if (end < 0)
1529 end = 0;
1530 }
1531
1532 /* zero-length substrings always match at the first attempt */
1533 if (pattern_len == 0)
1534 return (direction > 0) ? start : end;
1535
1536 end -= pattern_len;
1537
1538 if (direction < 0) {
1539 for (; end >= start; end--)
1540 if (Py_STRING_MATCH(target, end, pattern, pattern_len))
1541 return end;
1542 } else {
1543 for (; start <= end; start++)
1544 if (Py_STRING_MATCH(target, start, pattern, pattern_len))
1545 return start;
1546 }
1547 return -1;
1548}
1549
1550Py_LOCAL_INLINE(Py_ssize_t)
1551countstring(const char *target, Py_ssize_t target_len,
1552 const char *pattern, Py_ssize_t pattern_len,
1553 Py_ssize_t start,
1554 Py_ssize_t end,
1555 int direction, Py_ssize_t maxcount)
1556{
1557 Py_ssize_t count=0;
1558
1559 if (start < 0) {
1560 start += target_len;
1561 if (start < 0)
1562 start = 0;
1563 }
1564 if (end > target_len) {
1565 end = target_len;
1566 } else if (end < 0) {
1567 end += target_len;
1568 if (end < 0)
1569 end = 0;
1570 }
1571
1572 /* zero-length substrings match everywhere */
1573 if (pattern_len == 0 || maxcount == 0) {
1574 if (target_len+1 < maxcount)
1575 return target_len+1;
1576 return maxcount;
1577 }
1578
1579 end -= pattern_len;
1580 if (direction < 0) {
1581 for (; (end >= start); end--)
1582 if (Py_STRING_MATCH(target, end, pattern, pattern_len)) {
1583 count++;
1584 if (--maxcount <= 0) break;
1585 end -= pattern_len-1;
1586 }
1587 } else {
1588 for (; (start <= end); start++)
1589 if (Py_STRING_MATCH(target, start, pattern, pattern_len)) {
1590 count++;
1591 if (--maxcount <= 0)
1592 break;
1593 start += pattern_len-1;
1594 }
1595 }
1596 return count;
1597}
1598
1599
1600/* Algorithms for different cases of string replacement */
1601
1602/* len(self)>=1, from="", len(to)>=1, maxcount>=1 */
1603Py_LOCAL(PyByteArrayObject *)
1604replace_interleave(PyByteArrayObject *self,
1605 const char *to_s, Py_ssize_t to_len,
1606 Py_ssize_t maxcount)
1607{
1608 char *self_s, *result_s;
1609 Py_ssize_t self_len, result_len;
1610 Py_ssize_t count, i, product;
1611 PyByteArrayObject *result;
1612
1613 self_len = PyByteArray_GET_SIZE(self);
1614
1615 /* 1 at the end plus 1 after every character */
1616 count = self_len+1;
1617 if (maxcount < count)
1618 count = maxcount;
1619
1620 /* Check for overflow */
1621 /* result_len = count * to_len + self_len; */
1622 product = count * to_len;
1623 if (product / to_len != count) {
1624 PyErr_SetString(PyExc_OverflowError,
1625 "replace string is too long");
1626 return NULL;
1627 }
1628 result_len = product + self_len;
1629 if (result_len < 0) {
1630 PyErr_SetString(PyExc_OverflowError,
1631 "replace string is too long");
1632 return NULL;
1633 }
1634
1635 if (! (result = (PyByteArrayObject *)
1636 PyByteArray_FromStringAndSize(NULL, result_len)) )
1637 return NULL;
1638
1639 self_s = PyByteArray_AS_STRING(self);
1640 result_s = PyByteArray_AS_STRING(result);
1641
1642 /* TODO: special case single character, which doesn't need memcpy */
1643
1644 /* Lay the first one down (guaranteed this will occur) */
1645 Py_MEMCPY(result_s, to_s, to_len);
1646 result_s += to_len;
1647 count -= 1;
1648
1649 for (i=0; i<count; i++) {
1650 *result_s++ = *self_s++;
1651 Py_MEMCPY(result_s, to_s, to_len);
1652 result_s += to_len;
1653 }
1654
1655 /* Copy the rest of the original string */
1656 Py_MEMCPY(result_s, self_s, self_len-i);
1657
1658 return result;
1659}
1660
1661/* Special case for deleting a single character */
1662/* len(self)>=1, len(from)==1, to="", maxcount>=1 */
1663Py_LOCAL(PyByteArrayObject *)
1664replace_delete_single_character(PyByteArrayObject *self,
1665 char from_c, Py_ssize_t maxcount)
1666{
1667 char *self_s, *result_s;
1668 char *start, *next, *end;
1669 Py_ssize_t self_len, result_len;
1670 Py_ssize_t count;
1671 PyByteArrayObject *result;
1672
1673 self_len = PyByteArray_GET_SIZE(self);
1674 self_s = PyByteArray_AS_STRING(self);
1675
1676 count = countchar(self_s, self_len, from_c, maxcount);
1677 if (count == 0) {
1678 return return_self(self);
1679 }
1680
1681 result_len = self_len - count; /* from_len == 1 */
1682 assert(result_len>=0);
1683
1684 if ( (result = (PyByteArrayObject *)
1685 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL)
1686 return NULL;
1687 result_s = PyByteArray_AS_STRING(result);
1688
1689 start = self_s;
1690 end = self_s + self_len;
1691 while (count-- > 0) {
1692 next = findchar(start, end-start, from_c);
1693 if (next == NULL)
1694 break;
1695 Py_MEMCPY(result_s, start, next-start);
1696 result_s += (next-start);
1697 start = next+1;
1698 }
1699 Py_MEMCPY(result_s, start, end-start);
1700
1701 return result;
1702}
1703
1704/* len(self)>=1, len(from)>=2, to="", maxcount>=1 */
1705
1706Py_LOCAL(PyByteArrayObject *)
1707replace_delete_substring(PyByteArrayObject *self,
1708 const char *from_s, Py_ssize_t from_len,
1709 Py_ssize_t maxcount)
1710{
1711 char *self_s, *result_s;
1712 char *start, *next, *end;
1713 Py_ssize_t self_len, result_len;
1714 Py_ssize_t count, offset;
1715 PyByteArrayObject *result;
1716
1717 self_len = PyByteArray_GET_SIZE(self);
1718 self_s = PyByteArray_AS_STRING(self);
1719
1720 count = countstring(self_s, self_len,
1721 from_s, from_len,
1722 0, self_len, 1,
1723 maxcount);
1724
1725 if (count == 0) {
1726 /* no matches */
1727 return return_self(self);
1728 }
1729
1730 result_len = self_len - (count * from_len);
1731 assert (result_len>=0);
1732
1733 if ( (result = (PyByteArrayObject *)
1734 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL )
1735 return NULL;
1736
1737 result_s = PyByteArray_AS_STRING(result);
1738
1739 start = self_s;
1740 end = self_s + self_len;
1741 while (count-- > 0) {
1742 offset = findstring(start, end-start,
1743 from_s, from_len,
1744 0, end-start, FORWARD);
1745 if (offset == -1)
1746 break;
1747 next = start + offset;
1748
1749 Py_MEMCPY(result_s, start, next-start);
1750
1751 result_s += (next-start);
1752 start = next+from_len;
1753 }
1754 Py_MEMCPY(result_s, start, end-start);
1755 return result;
1756}
1757
1758/* len(self)>=1, len(from)==len(to)==1, maxcount>=1 */
1759Py_LOCAL(PyByteArrayObject *)
1760replace_single_character_in_place(PyByteArrayObject *self,
1761 char from_c, char to_c,
1762 Py_ssize_t maxcount)
1763{
Antoine Pitroua57aae72010-06-09 16:58:35 +00001764 char *self_s, *result_s, *start, *end, *next;
1765 Py_ssize_t self_len;
1766 PyByteArrayObject *result;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001767
Antoine Pitroua57aae72010-06-09 16:58:35 +00001768 /* The result string will be the same size */
1769 self_s = PyByteArray_AS_STRING(self);
1770 self_len = PyByteArray_GET_SIZE(self);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001771
Antoine Pitroua57aae72010-06-09 16:58:35 +00001772 next = findchar(self_s, self_len, from_c);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001773
Antoine Pitroua57aae72010-06-09 16:58:35 +00001774 if (next == NULL) {
1775 /* No matches; return the original bytes */
1776 return return_self(self);
1777 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001778
Antoine Pitroua57aae72010-06-09 16:58:35 +00001779 /* Need to make a new bytes */
1780 result = (PyByteArrayObject *) PyByteArray_FromStringAndSize(NULL, self_len);
1781 if (result == NULL)
1782 return NULL;
1783 result_s = PyByteArray_AS_STRING(result);
1784 Py_MEMCPY(result_s, self_s, self_len);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001785
Antoine Pitroua57aae72010-06-09 16:58:35 +00001786 /* change everything in-place, starting with this one */
1787 start = result_s + (next-self_s);
1788 *start = to_c;
1789 start++;
1790 end = result_s + self_len;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001791
Antoine Pitroua57aae72010-06-09 16:58:35 +00001792 while (--maxcount > 0) {
1793 next = findchar(start, end-start, from_c);
1794 if (next == NULL)
1795 break;
1796 *next = to_c;
1797 start = next+1;
1798 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001799
Antoine Pitroua57aae72010-06-09 16:58:35 +00001800 return result;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001801}
1802
1803/* len(self)>=1, len(from)==len(to)>=2, maxcount>=1 */
1804Py_LOCAL(PyByteArrayObject *)
1805replace_substring_in_place(PyByteArrayObject *self,
1806 const char *from_s, Py_ssize_t from_len,
1807 const char *to_s, Py_ssize_t to_len,
1808 Py_ssize_t maxcount)
1809{
1810 char *result_s, *start, *end;
1811 char *self_s;
1812 Py_ssize_t self_len, offset;
1813 PyByteArrayObject *result;
1814
1815 /* The result bytes will be the same size */
1816
1817 self_s = PyByteArray_AS_STRING(self);
1818 self_len = PyByteArray_GET_SIZE(self);
1819
1820 offset = findstring(self_s, self_len,
1821 from_s, from_len,
1822 0, self_len, FORWARD);
1823 if (offset == -1) {
1824 /* No matches; return the original bytes */
1825 return return_self(self);
1826 }
1827
1828 /* Need to make a new bytes */
1829 result = (PyByteArrayObject *) PyByteArray_FromStringAndSize(NULL, self_len);
1830 if (result == NULL)
1831 return NULL;
1832 result_s = PyByteArray_AS_STRING(result);
1833 Py_MEMCPY(result_s, self_s, self_len);
1834
1835 /* change everything in-place, starting with this one */
1836 start = result_s + offset;
1837 Py_MEMCPY(start, to_s, from_len);
1838 start += from_len;
1839 end = result_s + self_len;
1840
1841 while ( --maxcount > 0) {
1842 offset = findstring(start, end-start,
1843 from_s, from_len,
1844 0, end-start, FORWARD);
1845 if (offset==-1)
1846 break;
1847 Py_MEMCPY(start+offset, to_s, from_len);
1848 start += offset+from_len;
1849 }
1850
1851 return result;
1852}
1853
1854/* len(self)>=1, len(from)==1, len(to)>=2, maxcount>=1 */
1855Py_LOCAL(PyByteArrayObject *)
1856replace_single_character(PyByteArrayObject *self,
1857 char from_c,
1858 const char *to_s, Py_ssize_t to_len,
1859 Py_ssize_t maxcount)
1860{
1861 char *self_s, *result_s;
1862 char *start, *next, *end;
1863 Py_ssize_t self_len, result_len;
1864 Py_ssize_t count, product;
1865 PyByteArrayObject *result;
1866
1867 self_s = PyByteArray_AS_STRING(self);
1868 self_len = PyByteArray_GET_SIZE(self);
1869
1870 count = countchar(self_s, self_len, from_c, maxcount);
1871 if (count == 0) {
1872 /* no matches, return unchanged */
1873 return return_self(self);
1874 }
1875
1876 /* use the difference between current and new, hence the "-1" */
1877 /* result_len = self_len + count * (to_len-1) */
1878 product = count * (to_len-1);
1879 if (product / (to_len-1) != count) {
1880 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1881 return NULL;
1882 }
1883 result_len = self_len + product;
1884 if (result_len < 0) {
1885 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1886 return NULL;
1887 }
1888
1889 if ( (result = (PyByteArrayObject *)
1890 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL)
1891 return NULL;
1892 result_s = PyByteArray_AS_STRING(result);
1893
1894 start = self_s;
1895 end = self_s + self_len;
1896 while (count-- > 0) {
1897 next = findchar(start, end-start, from_c);
1898 if (next == NULL)
1899 break;
1900
1901 if (next == start) {
1902 /* replace with the 'to' */
1903 Py_MEMCPY(result_s, to_s, to_len);
1904 result_s += to_len;
1905 start += 1;
1906 } else {
1907 /* copy the unchanged old then the 'to' */
1908 Py_MEMCPY(result_s, start, next-start);
1909 result_s += (next-start);
1910 Py_MEMCPY(result_s, to_s, to_len);
1911 result_s += to_len;
1912 start = next+1;
1913 }
1914 }
1915 /* Copy the remainder of the remaining bytes */
1916 Py_MEMCPY(result_s, start, end-start);
1917
1918 return result;
1919}
1920
1921/* len(self)>=1, len(from)>=2, len(to)>=2, maxcount>=1 */
1922Py_LOCAL(PyByteArrayObject *)
1923replace_substring(PyByteArrayObject *self,
1924 const char *from_s, Py_ssize_t from_len,
1925 const char *to_s, Py_ssize_t to_len,
1926 Py_ssize_t maxcount)
1927{
1928 char *self_s, *result_s;
1929 char *start, *next, *end;
1930 Py_ssize_t self_len, result_len;
1931 Py_ssize_t count, offset, product;
1932 PyByteArrayObject *result;
1933
1934 self_s = PyByteArray_AS_STRING(self);
1935 self_len = PyByteArray_GET_SIZE(self);
1936
1937 count = countstring(self_s, self_len,
1938 from_s, from_len,
1939 0, self_len, FORWARD, maxcount);
1940 if (count == 0) {
1941 /* no matches, return unchanged */
1942 return return_self(self);
1943 }
1944
1945 /* Check for overflow */
1946 /* result_len = self_len + count * (to_len-from_len) */
1947 product = count * (to_len-from_len);
1948 if (product / (to_len-from_len) != count) {
1949 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1950 return NULL;
1951 }
1952 result_len = self_len + product;
1953 if (result_len < 0) {
1954 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1955 return NULL;
1956 }
1957
1958 if ( (result = (PyByteArrayObject *)
1959 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL)
1960 return NULL;
1961 result_s = PyByteArray_AS_STRING(result);
1962
1963 start = self_s;
1964 end = self_s + self_len;
1965 while (count-- > 0) {
1966 offset = findstring(start, end-start,
1967 from_s, from_len,
1968 0, end-start, FORWARD);
1969 if (offset == -1)
1970 break;
1971 next = start+offset;
1972 if (next == start) {
1973 /* replace with the 'to' */
1974 Py_MEMCPY(result_s, to_s, to_len);
1975 result_s += to_len;
1976 start += from_len;
1977 } else {
1978 /* copy the unchanged old then the 'to' */
1979 Py_MEMCPY(result_s, start, next-start);
1980 result_s += (next-start);
1981 Py_MEMCPY(result_s, to_s, to_len);
1982 result_s += to_len;
1983 start = next+from_len;
1984 }
1985 }
1986 /* Copy the remainder of the remaining bytes */
1987 Py_MEMCPY(result_s, start, end-start);
1988
1989 return result;
1990}
1991
1992
1993Py_LOCAL(PyByteArrayObject *)
1994replace(PyByteArrayObject *self,
1995 const char *from_s, Py_ssize_t from_len,
1996 const char *to_s, Py_ssize_t to_len,
1997 Py_ssize_t maxcount)
1998{
1999 if (maxcount < 0) {
2000 maxcount = PY_SSIZE_T_MAX;
2001 } else if (maxcount == 0 || PyByteArray_GET_SIZE(self) == 0) {
2002 /* nothing to do; return the original bytes */
2003 return return_self(self);
2004 }
2005
2006 if (maxcount == 0 ||
2007 (from_len == 0 && to_len == 0)) {
2008 /* nothing to do; return the original bytes */
2009 return return_self(self);
2010 }
2011
2012 /* Handle zero-length special cases */
2013
2014 if (from_len == 0) {
2015 /* insert the 'to' bytes everywhere. */
2016 /* >>> "Python".replace("", ".") */
2017 /* '.P.y.t.h.o.n.' */
2018 return replace_interleave(self, to_s, to_len, maxcount);
2019 }
2020
2021 /* Except for "".replace("", "A") == "A" there is no way beyond this */
2022 /* point for an empty self bytes to generate a non-empty bytes */
2023 /* Special case so the remaining code always gets a non-empty bytes */
2024 if (PyByteArray_GET_SIZE(self) == 0) {
2025 return return_self(self);
2026 }
2027
2028 if (to_len == 0) {
Georg Brandl17cb8a82008-05-30 08:20:09 +00002029 /* delete all occurrences of 'from' bytes */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002030 if (from_len == 1) {
2031 return replace_delete_single_character(
2032 self, from_s[0], maxcount);
2033 } else {
2034 return replace_delete_substring(self, from_s, from_len, maxcount);
2035 }
2036 }
2037
2038 /* Handle special case where both bytes have the same length */
2039
2040 if (from_len == to_len) {
2041 if (from_len == 1) {
2042 return replace_single_character_in_place(
2043 self,
2044 from_s[0],
2045 to_s[0],
2046 maxcount);
2047 } else {
2048 return replace_substring_in_place(
2049 self, from_s, from_len, to_s, to_len, maxcount);
2050 }
2051 }
2052
2053 /* Otherwise use the more generic algorithms */
2054 if (from_len == 1) {
2055 return replace_single_character(self, from_s[0],
2056 to_s, to_len, maxcount);
2057 } else {
2058 /* len('from')>=2, len('to')>=1 */
2059 return replace_substring(self, from_s, from_len, to_s, to_len, maxcount);
2060 }
2061}
2062
2063
2064PyDoc_STRVAR(replace__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002065"B.replace(old, new[, count]) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002066\n\
2067Return a copy of B with all occurrences of subsection\n\
2068old replaced by new. If the optional argument count is\n\
2069given, only the first count occurrences are replaced.");
2070
2071static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002072bytearray_replace(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002073{
2074 Py_ssize_t count = -1;
2075 PyObject *from, *to, *res;
2076 Py_buffer vfrom, vto;
2077
2078 if (!PyArg_ParseTuple(args, "OO|n:replace", &from, &to, &count))
2079 return NULL;
2080
2081 if (_getbuffer(from, &vfrom) < 0)
2082 return NULL;
2083 if (_getbuffer(to, &vto) < 0) {
Martin v. Löwis423be952008-08-13 15:53:07 +00002084 PyBuffer_Release(&vfrom);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002085 return NULL;
2086 }
2087
2088 res = (PyObject *)replace((PyByteArrayObject *) self,
2089 vfrom.buf, vfrom.len,
2090 vto.buf, vto.len, count);
2091
Martin v. Löwis423be952008-08-13 15:53:07 +00002092 PyBuffer_Release(&vfrom);
2093 PyBuffer_Release(&vto);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002094 return res;
2095}
2096
2097
2098/* Overallocate the initial list to reduce the number of reallocs for small
2099 split sizes. Eg, "A A A A A A A A A A".split() (10 elements) has three
2100 resizes, to sizes 4, 8, then 16. Most observed string splits are for human
2101 text (roughly 11 words per line) and field delimited data (usually 1-10
2102 fields). For large strings the split algorithms are bandwidth limited
2103 so increasing the preallocation likely will not improve things.*/
2104
2105#define MAX_PREALLOC 12
2106
2107/* 5 splits gives 6 elements */
2108#define PREALLOC_SIZE(maxsplit) \
2109 (maxsplit >= MAX_PREALLOC ? MAX_PREALLOC : maxsplit+1)
2110
2111#define SPLIT_APPEND(data, left, right) \
2112 str = PyByteArray_FromStringAndSize((data) + (left), \
2113 (right) - (left)); \
2114 if (str == NULL) \
2115 goto onError; \
2116 if (PyList_Append(list, str)) { \
2117 Py_DECREF(str); \
2118 goto onError; \
2119 } \
2120 else \
2121 Py_DECREF(str);
2122
2123#define SPLIT_ADD(data, left, right) { \
2124 str = PyByteArray_FromStringAndSize((data) + (left), \
2125 (right) - (left)); \
2126 if (str == NULL) \
2127 goto onError; \
2128 if (count < MAX_PREALLOC) { \
2129 PyList_SET_ITEM(list, count, str); \
2130 } else { \
2131 if (PyList_Append(list, str)) { \
2132 Py_DECREF(str); \
2133 goto onError; \
2134 } \
2135 else \
2136 Py_DECREF(str); \
2137 } \
2138 count++; }
2139
2140/* Always force the list to the expected size. */
2141#define FIX_PREALLOC_SIZE(list) Py_SIZE(list) = count
2142
2143
2144Py_LOCAL_INLINE(PyObject *)
2145split_char(const char *s, Py_ssize_t len, char ch, Py_ssize_t maxcount)
2146{
2147 register Py_ssize_t i, j, count = 0;
2148 PyObject *str;
2149 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2150
2151 if (list == NULL)
2152 return NULL;
2153
2154 i = j = 0;
2155 while ((j < len) && (maxcount-- > 0)) {
2156 for(; j < len; j++) {
2157 /* I found that using memchr makes no difference */
2158 if (s[j] == ch) {
2159 SPLIT_ADD(s, i, j);
2160 i = j = j + 1;
2161 break;
2162 }
2163 }
2164 }
2165 if (i <= len) {
2166 SPLIT_ADD(s, i, len);
2167 }
2168 FIX_PREALLOC_SIZE(list);
2169 return list;
2170
2171 onError:
2172 Py_DECREF(list);
2173 return NULL;
2174}
2175
2176
2177Py_LOCAL_INLINE(PyObject *)
2178split_whitespace(const char *s, Py_ssize_t len, Py_ssize_t maxcount)
2179{
2180 register Py_ssize_t i, j, count = 0;
2181 PyObject *str;
2182 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2183
2184 if (list == NULL)
2185 return NULL;
2186
2187 for (i = j = 0; i < len; ) {
2188 /* find a token */
Eric Smith6dc46f52009-04-27 20:39:49 +00002189 while (i < len && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002190 i++;
2191 j = i;
Eric Smith6dc46f52009-04-27 20:39:49 +00002192 while (i < len && !Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002193 i++;
2194 if (j < i) {
2195 if (maxcount-- <= 0)
2196 break;
2197 SPLIT_ADD(s, j, i);
Eric Smith6dc46f52009-04-27 20:39:49 +00002198 while (i < len && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002199 i++;
2200 j = i;
2201 }
2202 }
2203 if (j < len) {
2204 SPLIT_ADD(s, j, len);
2205 }
2206 FIX_PREALLOC_SIZE(list);
2207 return list;
2208
2209 onError:
2210 Py_DECREF(list);
2211 return NULL;
2212}
2213
2214PyDoc_STRVAR(split__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002215"B.split([sep[, maxsplit]]) -> list of bytearrays\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002216\n\
2217Return a list of the sections in B, using sep as the delimiter.\n\
2218If sep is not given, B is split on ASCII whitespace characters\n\
2219(space, tab, return, newline, formfeed, vertical tab).\n\
2220If maxsplit is given, at most maxsplit splits are done.");
2221
2222static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002223bytearray_split(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002224{
2225 Py_ssize_t len = PyByteArray_GET_SIZE(self), n, i, j;
2226 Py_ssize_t maxsplit = -1, count = 0;
2227 const char *s = PyByteArray_AS_STRING(self), *sub;
2228 PyObject *list, *str, *subobj = Py_None;
2229 Py_buffer vsub;
2230#ifdef USE_FAST
2231 Py_ssize_t pos;
2232#endif
2233
2234 if (!PyArg_ParseTuple(args, "|On:split", &subobj, &maxsplit))
2235 return NULL;
2236 if (maxsplit < 0)
2237 maxsplit = PY_SSIZE_T_MAX;
2238
2239 if (subobj == Py_None)
2240 return split_whitespace(s, len, maxsplit);
2241
2242 if (_getbuffer(subobj, &vsub) < 0)
2243 return NULL;
2244 sub = vsub.buf;
2245 n = vsub.len;
2246
2247 if (n == 0) {
2248 PyErr_SetString(PyExc_ValueError, "empty separator");
Martin v. Löwis423be952008-08-13 15:53:07 +00002249 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002250 return NULL;
2251 }
Benjamin Petersonc4fe6f32008-08-19 18:57:56 +00002252 if (n == 1) {
2253 list = split_char(s, len, sub[0], maxsplit);
2254 PyBuffer_Release(&vsub);
2255 return list;
2256 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002257
2258 list = PyList_New(PREALLOC_SIZE(maxsplit));
2259 if (list == NULL) {
Martin v. Löwis423be952008-08-13 15:53:07 +00002260 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002261 return NULL;
2262 }
2263
2264#ifdef USE_FAST
2265 i = j = 0;
2266 while (maxsplit-- > 0) {
2267 pos = fastsearch(s+i, len-i, sub, n, FAST_SEARCH);
2268 if (pos < 0)
2269 break;
2270 j = i+pos;
2271 SPLIT_ADD(s, i, j);
2272 i = j + n;
2273 }
2274#else
2275 i = j = 0;
2276 while ((j+n <= len) && (maxsplit-- > 0)) {
2277 for (; j+n <= len; j++) {
2278 if (Py_STRING_MATCH(s, j, sub, n)) {
2279 SPLIT_ADD(s, i, j);
2280 i = j = j + n;
2281 break;
2282 }
2283 }
2284 }
2285#endif
2286 SPLIT_ADD(s, i, len);
2287 FIX_PREALLOC_SIZE(list);
Martin v. Löwis423be952008-08-13 15:53:07 +00002288 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002289 return list;
2290
2291 onError:
2292 Py_DECREF(list);
Martin v. Löwis423be952008-08-13 15:53:07 +00002293 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002294 return NULL;
2295}
2296
2297/* stringlib's partition shares nullbytes in some cases.
2298 undo this, we don't want the nullbytes to be shared. */
2299static PyObject *
2300make_nullbytes_unique(PyObject *result)
2301{
2302 if (result != NULL) {
2303 int i;
2304 assert(PyTuple_Check(result));
2305 assert(PyTuple_GET_SIZE(result) == 3);
2306 for (i = 0; i < 3; i++) {
2307 if (PyTuple_GET_ITEM(result, i) == (PyObject *)nullbytes) {
2308 PyObject *new = PyByteArray_FromStringAndSize(NULL, 0);
2309 if (new == NULL) {
2310 Py_DECREF(result);
2311 result = NULL;
2312 break;
2313 }
2314 Py_DECREF(nullbytes);
2315 PyTuple_SET_ITEM(result, i, new);
2316 }
2317 }
2318 }
2319 return result;
2320}
2321
2322PyDoc_STRVAR(partition__doc__,
2323"B.partition(sep) -> (head, sep, tail)\n\
2324\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002325Search for the separator sep in B, and return the part before it,\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002326the separator itself, and the part after it. If the separator is not\n\
2327found, returns B and two empty bytearray objects.");
2328
2329static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002330bytearray_partition(PyByteArrayObject *self, PyObject *sep_obj)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002331{
2332 PyObject *bytesep, *result;
2333
2334 bytesep = PyByteArray_FromObject(sep_obj);
2335 if (! bytesep)
2336 return NULL;
2337
2338 result = stringlib_partition(
2339 (PyObject*) self,
2340 PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
2341 bytesep,
2342 PyByteArray_AS_STRING(bytesep), PyByteArray_GET_SIZE(bytesep)
2343 );
2344
2345 Py_DECREF(bytesep);
2346 return make_nullbytes_unique(result);
2347}
2348
2349PyDoc_STRVAR(rpartition__doc__,
Ezio Melotti4c81fbb2010-01-25 12:02:24 +00002350"B.rpartition(sep) -> (head, sep, tail)\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002351\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002352Search for the separator sep in B, starting at the end of B,\n\
2353and return the part before it, the separator itself, and the\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002354part after it. If the separator is not found, returns two empty\n\
2355bytearray objects and B.");
2356
2357static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002358bytearray_rpartition(PyByteArrayObject *self, PyObject *sep_obj)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002359{
2360 PyObject *bytesep, *result;
2361
2362 bytesep = PyByteArray_FromObject(sep_obj);
2363 if (! bytesep)
2364 return NULL;
2365
2366 result = stringlib_rpartition(
2367 (PyObject*) self,
2368 PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
2369 bytesep,
2370 PyByteArray_AS_STRING(bytesep), PyByteArray_GET_SIZE(bytesep)
2371 );
2372
2373 Py_DECREF(bytesep);
2374 return make_nullbytes_unique(result);
2375}
2376
2377Py_LOCAL_INLINE(PyObject *)
2378rsplit_char(const char *s, Py_ssize_t len, char ch, Py_ssize_t maxcount)
2379{
2380 register Py_ssize_t i, j, count=0;
2381 PyObject *str;
2382 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2383
2384 if (list == NULL)
2385 return NULL;
2386
2387 i = j = len - 1;
2388 while ((i >= 0) && (maxcount-- > 0)) {
2389 for (; i >= 0; i--) {
2390 if (s[i] == ch) {
2391 SPLIT_ADD(s, i + 1, j + 1);
2392 j = i = i - 1;
2393 break;
2394 }
2395 }
2396 }
2397 if (j >= -1) {
2398 SPLIT_ADD(s, 0, j + 1);
2399 }
2400 FIX_PREALLOC_SIZE(list);
2401 if (PyList_Reverse(list) < 0)
2402 goto onError;
2403
2404 return list;
2405
2406 onError:
2407 Py_DECREF(list);
2408 return NULL;
2409}
2410
2411Py_LOCAL_INLINE(PyObject *)
2412rsplit_whitespace(const char *s, Py_ssize_t len, Py_ssize_t maxcount)
2413{
2414 register Py_ssize_t i, j, count = 0;
2415 PyObject *str;
2416 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2417
2418 if (list == NULL)
2419 return NULL;
2420
2421 for (i = j = len - 1; i >= 0; ) {
2422 /* find a token */
Eric Smith6dc46f52009-04-27 20:39:49 +00002423 while (i >= 0 && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002424 i--;
2425 j = i;
Eric Smith6dc46f52009-04-27 20:39:49 +00002426 while (i >= 0 && !Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002427 i--;
2428 if (j > i) {
2429 if (maxcount-- <= 0)
2430 break;
2431 SPLIT_ADD(s, i + 1, j + 1);
Eric Smith6dc46f52009-04-27 20:39:49 +00002432 while (i >= 0 && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002433 i--;
2434 j = i;
2435 }
2436 }
2437 if (j >= 0) {
2438 SPLIT_ADD(s, 0, j + 1);
2439 }
2440 FIX_PREALLOC_SIZE(list);
2441 if (PyList_Reverse(list) < 0)
2442 goto onError;
2443
2444 return list;
2445
2446 onError:
2447 Py_DECREF(list);
2448 return NULL;
2449}
2450
2451PyDoc_STRVAR(rsplit__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002452"B.rsplit(sep[, maxsplit]) -> list of bytearrays\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002453\n\
2454Return a list of the sections in B, using sep as the delimiter,\n\
2455starting at the end of B and working to the front.\n\
2456If sep is not given, B is split on ASCII whitespace characters\n\
2457(space, tab, return, newline, formfeed, vertical tab).\n\
2458If maxsplit is given, at most maxsplit splits are done.");
2459
2460static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002461bytearray_rsplit(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002462{
2463 Py_ssize_t len = PyByteArray_GET_SIZE(self), n, i, j;
2464 Py_ssize_t maxsplit = -1, count = 0;
2465 const char *s = PyByteArray_AS_STRING(self), *sub;
2466 PyObject *list, *str, *subobj = Py_None;
2467 Py_buffer vsub;
2468
2469 if (!PyArg_ParseTuple(args, "|On:rsplit", &subobj, &maxsplit))
2470 return NULL;
2471 if (maxsplit < 0)
2472 maxsplit = PY_SSIZE_T_MAX;
2473
2474 if (subobj == Py_None)
2475 return rsplit_whitespace(s, len, maxsplit);
2476
2477 if (_getbuffer(subobj, &vsub) < 0)
2478 return NULL;
2479 sub = vsub.buf;
2480 n = vsub.len;
2481
2482 if (n == 0) {
2483 PyErr_SetString(PyExc_ValueError, "empty separator");
Martin v. Löwis423be952008-08-13 15:53:07 +00002484 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002485 return NULL;
2486 }
Benjamin Petersonc4fe6f32008-08-19 18:57:56 +00002487 else if (n == 1) {
2488 list = rsplit_char(s, len, sub[0], maxsplit);
2489 PyBuffer_Release(&vsub);
2490 return list;
2491 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002492
2493 list = PyList_New(PREALLOC_SIZE(maxsplit));
2494 if (list == NULL) {
Martin v. Löwis423be952008-08-13 15:53:07 +00002495 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002496 return NULL;
2497 }
2498
2499 j = len;
2500 i = j - n;
2501
2502 while ( (i >= 0) && (maxsplit-- > 0) ) {
2503 for (; i>=0; i--) {
2504 if (Py_STRING_MATCH(s, i, sub, n)) {
2505 SPLIT_ADD(s, i + n, j);
2506 j = i;
2507 i -= n;
2508 break;
2509 }
2510 }
2511 }
2512 SPLIT_ADD(s, 0, j);
2513 FIX_PREALLOC_SIZE(list);
2514 if (PyList_Reverse(list) < 0)
2515 goto onError;
Martin v. Löwis423be952008-08-13 15:53:07 +00002516 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002517 return list;
2518
2519onError:
2520 Py_DECREF(list);
Martin v. Löwis423be952008-08-13 15:53:07 +00002521 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002522 return NULL;
2523}
2524
2525PyDoc_STRVAR(reverse__doc__,
2526"B.reverse() -> None\n\
2527\n\
2528Reverse the order of the values in B in place.");
2529static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002530bytearray_reverse(PyByteArrayObject *self, PyObject *unused)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002531{
2532 char swap, *head, *tail;
2533 Py_ssize_t i, j, n = Py_SIZE(self);
2534
2535 j = n / 2;
2536 head = self->ob_bytes;
2537 tail = head + n - 1;
2538 for (i = 0; i < j; i++) {
2539 swap = *head;
2540 *head++ = *tail;
2541 *tail-- = swap;
2542 }
2543
2544 Py_RETURN_NONE;
2545}
2546
2547PyDoc_STRVAR(insert__doc__,
2548"B.insert(index, int) -> None\n\
2549\n\
2550Insert a single item into the bytearray before the given index.");
2551static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002552bytearray_insert(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002553{
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002554 PyObject *value;
2555 int ival;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002556 Py_ssize_t where, n = Py_SIZE(self);
2557
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002558 if (!PyArg_ParseTuple(args, "nO:insert", &where, &value))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002559 return NULL;
2560
2561 if (n == PY_SSIZE_T_MAX) {
2562 PyErr_SetString(PyExc_OverflowError,
Mark Dickinson099d7a02009-09-06 10:35:38 +00002563 "cannot add more objects to bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002564 return NULL;
2565 }
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002566 if (!_getbytevalue(value, &ival))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002567 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002568 if (PyByteArray_Resize((PyObject *)self, n + 1) < 0)
2569 return NULL;
2570
2571 if (where < 0) {
2572 where += n;
2573 if (where < 0)
2574 where = 0;
2575 }
2576 if (where > n)
2577 where = n;
2578 memmove(self->ob_bytes + where + 1, self->ob_bytes + where, n - where);
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002579 self->ob_bytes[where] = ival;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002580
2581 Py_RETURN_NONE;
2582}
2583
2584PyDoc_STRVAR(append__doc__,
2585"B.append(int) -> None\n\
2586\n\
2587Append a single item to the end of B.");
2588static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002589bytearray_append(PyByteArrayObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002590{
2591 int value;
2592 Py_ssize_t n = Py_SIZE(self);
2593
2594 if (! _getbytevalue(arg, &value))
2595 return NULL;
2596 if (n == PY_SSIZE_T_MAX) {
2597 PyErr_SetString(PyExc_OverflowError,
Mark Dickinson099d7a02009-09-06 10:35:38 +00002598 "cannot add more objects to bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002599 return NULL;
2600 }
2601 if (PyByteArray_Resize((PyObject *)self, n + 1) < 0)
2602 return NULL;
2603
2604 self->ob_bytes[n] = value;
2605
2606 Py_RETURN_NONE;
2607}
2608
2609PyDoc_STRVAR(extend__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002610"B.extend(iterable_of_ints) -> None\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002611\n\
2612Append all the elements from the iterator or sequence to the\n\
2613end of B.");
2614static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002615bytearray_extend(PyByteArrayObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002616{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002617 PyObject *it, *item, *bytearray_obj;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002618 Py_ssize_t buf_size = 0, len = 0;
2619 int value;
2620 char *buf;
2621
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002622 /* bytearray_setslice code only accepts something supporting PEP 3118. */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002623 if (PyObject_CheckBuffer(arg)) {
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002624 if (bytearray_setslice(self, Py_SIZE(self), Py_SIZE(self), arg) == -1)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002625 return NULL;
2626
2627 Py_RETURN_NONE;
2628 }
2629
2630 it = PyObject_GetIter(arg);
2631 if (it == NULL)
2632 return NULL;
2633
Ezio Melotti42da6632011-03-15 05:18:48 +02002634 /* Try to determine the length of the argument. 32 is arbitrary. */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002635 buf_size = _PyObject_LengthHint(arg, 32);
Benjamin Petersonef3e4c22009-04-11 19:48:14 +00002636 if (buf_size == -1) {
2637 Py_DECREF(it);
2638 return NULL;
2639 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002640
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002641 bytearray_obj = PyByteArray_FromStringAndSize(NULL, buf_size);
2642 if (bytearray_obj == NULL)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002643 return NULL;
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002644 buf = PyByteArray_AS_STRING(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002645
2646 while ((item = PyIter_Next(it)) != NULL) {
2647 if (! _getbytevalue(item, &value)) {
2648 Py_DECREF(item);
2649 Py_DECREF(it);
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002650 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002651 return NULL;
2652 }
2653 buf[len++] = value;
2654 Py_DECREF(item);
2655
2656 if (len >= buf_size) {
2657 buf_size = len + (len >> 1) + 1;
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002658 if (PyByteArray_Resize((PyObject *)bytearray_obj, buf_size) < 0) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002659 Py_DECREF(it);
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002660 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002661 return NULL;
2662 }
2663 /* Recompute the `buf' pointer, since the resizing operation may
2664 have invalidated it. */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002665 buf = PyByteArray_AS_STRING(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002666 }
2667 }
2668 Py_DECREF(it);
2669
2670 /* Resize down to exact size. */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002671 if (PyByteArray_Resize((PyObject *)bytearray_obj, len) < 0) {
2672 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002673 return NULL;
2674 }
2675
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002676 if (bytearray_setslice(self, Py_SIZE(self), Py_SIZE(self), bytearray_obj) == -1)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002677 return NULL;
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002678 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002679
2680 Py_RETURN_NONE;
2681}
2682
2683PyDoc_STRVAR(pop__doc__,
2684"B.pop([index]) -> int\n\
2685\n\
2686Remove and return a single item from B. If no index\n\
Benjamin Petersondcf97b92008-07-02 17:30:14 +00002687argument is given, will pop the last value.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002688static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002689bytearray_pop(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002690{
2691 int value;
2692 Py_ssize_t where = -1, n = Py_SIZE(self);
2693
2694 if (!PyArg_ParseTuple(args, "|n:pop", &where))
2695 return NULL;
2696
2697 if (n == 0) {
2698 PyErr_SetString(PyExc_OverflowError,
Mark Dickinson099d7a02009-09-06 10:35:38 +00002699 "cannot pop an empty bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002700 return NULL;
2701 }
2702 if (where < 0)
2703 where += Py_SIZE(self);
2704 if (where < 0 || where >= Py_SIZE(self)) {
2705 PyErr_SetString(PyExc_IndexError, "pop index out of range");
2706 return NULL;
2707 }
Antoine Pitrou5504e892008-12-06 21:27:53 +00002708 if (!_canresize(self))
2709 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002710
2711 value = self->ob_bytes[where];
2712 memmove(self->ob_bytes + where, self->ob_bytes + where + 1, n - where);
2713 if (PyByteArray_Resize((PyObject *)self, n - 1) < 0)
2714 return NULL;
2715
Mark Dickinson424d75a2009-09-06 10:20:23 +00002716 return PyLong_FromLong((unsigned char)value);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002717}
2718
2719PyDoc_STRVAR(remove__doc__,
2720"B.remove(int) -> None\n\
2721\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002722Remove the first occurrence of a value in B.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002723static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002724bytearray_remove(PyByteArrayObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002725{
2726 int value;
2727 Py_ssize_t where, n = Py_SIZE(self);
2728
2729 if (! _getbytevalue(arg, &value))
2730 return NULL;
2731
2732 for (where = 0; where < n; where++) {
2733 if (self->ob_bytes[where] == value)
2734 break;
2735 }
2736 if (where == n) {
Mark Dickinson099d7a02009-09-06 10:35:38 +00002737 PyErr_SetString(PyExc_ValueError, "value not found in bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002738 return NULL;
2739 }
Antoine Pitrou5504e892008-12-06 21:27:53 +00002740 if (!_canresize(self))
2741 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002742
2743 memmove(self->ob_bytes + where, self->ob_bytes + where + 1, n - where);
2744 if (PyByteArray_Resize((PyObject *)self, n - 1) < 0)
2745 return NULL;
2746
2747 Py_RETURN_NONE;
2748}
2749
2750/* XXX These two helpers could be optimized if argsize == 1 */
2751
2752static Py_ssize_t
2753lstrip_helper(unsigned char *myptr, Py_ssize_t mysize,
2754 void *argptr, Py_ssize_t argsize)
2755{
2756 Py_ssize_t i = 0;
2757 while (i < mysize && memchr(argptr, myptr[i], argsize))
2758 i++;
2759 return i;
2760}
2761
2762static Py_ssize_t
2763rstrip_helper(unsigned char *myptr, Py_ssize_t mysize,
2764 void *argptr, Py_ssize_t argsize)
2765{
2766 Py_ssize_t i = mysize - 1;
2767 while (i >= 0 && memchr(argptr, myptr[i], argsize))
2768 i--;
2769 return i + 1;
2770}
2771
2772PyDoc_STRVAR(strip__doc__,
2773"B.strip([bytes]) -> bytearray\n\
2774\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002775Strip leading and trailing bytes contained in the argument\n\
2776and return the result as a new bytearray.\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002777If the argument is omitted, strip ASCII whitespace.");
2778static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002779bytearray_strip(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002780{
2781 Py_ssize_t left, right, mysize, argsize;
2782 void *myptr, *argptr;
2783 PyObject *arg = Py_None;
2784 Py_buffer varg;
2785 if (!PyArg_ParseTuple(args, "|O:strip", &arg))
2786 return NULL;
2787 if (arg == Py_None) {
2788 argptr = "\t\n\r\f\v ";
2789 argsize = 6;
2790 }
2791 else {
2792 if (_getbuffer(arg, &varg) < 0)
2793 return NULL;
2794 argptr = varg.buf;
2795 argsize = varg.len;
2796 }
2797 myptr = self->ob_bytes;
2798 mysize = Py_SIZE(self);
2799 left = lstrip_helper(myptr, mysize, argptr, argsize);
2800 if (left == mysize)
2801 right = left;
2802 else
2803 right = rstrip_helper(myptr, mysize, argptr, argsize);
2804 if (arg != Py_None)
Martin v. Löwis423be952008-08-13 15:53:07 +00002805 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002806 return PyByteArray_FromStringAndSize(self->ob_bytes + left, right - left);
2807}
2808
2809PyDoc_STRVAR(lstrip__doc__,
2810"B.lstrip([bytes]) -> bytearray\n\
2811\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002812Strip leading bytes contained in the argument\n\
2813and return the result as a new bytearray.\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002814If the argument is omitted, strip leading ASCII whitespace.");
2815static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002816bytearray_lstrip(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002817{
2818 Py_ssize_t left, right, mysize, argsize;
2819 void *myptr, *argptr;
2820 PyObject *arg = Py_None;
2821 Py_buffer varg;
2822 if (!PyArg_ParseTuple(args, "|O:lstrip", &arg))
2823 return NULL;
2824 if (arg == Py_None) {
2825 argptr = "\t\n\r\f\v ";
2826 argsize = 6;
2827 }
2828 else {
2829 if (_getbuffer(arg, &varg) < 0)
2830 return NULL;
2831 argptr = varg.buf;
2832 argsize = varg.len;
2833 }
2834 myptr = self->ob_bytes;
2835 mysize = Py_SIZE(self);
2836 left = lstrip_helper(myptr, mysize, argptr, argsize);
2837 right = mysize;
2838 if (arg != Py_None)
Martin v. Löwis423be952008-08-13 15:53:07 +00002839 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002840 return PyByteArray_FromStringAndSize(self->ob_bytes + left, right - left);
2841}
2842
2843PyDoc_STRVAR(rstrip__doc__,
2844"B.rstrip([bytes]) -> bytearray\n\
2845\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002846Strip trailing bytes contained in the argument\n\
2847and return the result as a new bytearray.\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002848If the argument is omitted, strip trailing ASCII whitespace.");
2849static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002850bytearray_rstrip(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002851{
2852 Py_ssize_t left, right, mysize, argsize;
2853 void *myptr, *argptr;
2854 PyObject *arg = Py_None;
2855 Py_buffer varg;
2856 if (!PyArg_ParseTuple(args, "|O:rstrip", &arg))
2857 return NULL;
2858 if (arg == Py_None) {
2859 argptr = "\t\n\r\f\v ";
2860 argsize = 6;
2861 }
2862 else {
2863 if (_getbuffer(arg, &varg) < 0)
2864 return NULL;
2865 argptr = varg.buf;
2866 argsize = varg.len;
2867 }
2868 myptr = self->ob_bytes;
2869 mysize = Py_SIZE(self);
2870 left = 0;
2871 right = rstrip_helper(myptr, mysize, argptr, argsize);
2872 if (arg != Py_None)
Martin v. Löwis423be952008-08-13 15:53:07 +00002873 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002874 return PyByteArray_FromStringAndSize(self->ob_bytes + left, right - left);
2875}
2876
2877PyDoc_STRVAR(decode_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002878"B.decode([encoding[, errors]]) -> str\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002879\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002880Decode B using the codec registered for encoding. encoding defaults\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002881to the default encoding. errors may be given to set a different error\n\
2882handling scheme. Default is 'strict' meaning that encoding errors raise\n\
2883a UnicodeDecodeError. Other possible values are 'ignore' and 'replace'\n\
2884as well as any other name registered with codecs.register_error that is\n\
2885able to handle UnicodeDecodeErrors.");
2886
2887static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002888bytearray_decode(PyObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002889{
2890 const char *encoding = NULL;
2891 const char *errors = NULL;
2892
2893 if (!PyArg_ParseTuple(args, "|ss:decode", &encoding, &errors))
2894 return NULL;
2895 if (encoding == NULL)
2896 encoding = PyUnicode_GetDefaultEncoding();
Marc-André Lemburgb2750b52008-06-06 12:18:17 +00002897 return PyUnicode_FromEncodedObject(self, encoding, errors);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002898}
2899
2900PyDoc_STRVAR(alloc_doc,
2901"B.__alloc__() -> int\n\
2902\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002903Return the number of bytes actually allocated.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002904
2905static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002906bytearray_alloc(PyByteArrayObject *self)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002907{
2908 return PyLong_FromSsize_t(self->ob_alloc);
2909}
2910
2911PyDoc_STRVAR(join_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002912"B.join(iterable_of_bytes) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002913\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002914Concatenate any number of bytes/bytearray objects, with B\n\
2915in between each pair, and return the result as a new bytearray.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002916
2917static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002918bytearray_join(PyByteArrayObject *self, PyObject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002919{
2920 PyObject *seq;
2921 Py_ssize_t mysize = Py_SIZE(self);
2922 Py_ssize_t i;
2923 Py_ssize_t n;
2924 PyObject **items;
2925 Py_ssize_t totalsize = 0;
2926 PyObject *result;
2927 char *dest;
2928
2929 seq = PySequence_Fast(it, "can only join an iterable");
2930 if (seq == NULL)
2931 return NULL;
2932 n = PySequence_Fast_GET_SIZE(seq);
2933 items = PySequence_Fast_ITEMS(seq);
2934
2935 /* Compute the total size, and check that they are all bytes */
2936 /* XXX Shouldn't we use _getbuffer() on these items instead? */
2937 for (i = 0; i < n; i++) {
2938 PyObject *obj = items[i];
2939 if (!PyByteArray_Check(obj) && !PyBytes_Check(obj)) {
2940 PyErr_Format(PyExc_TypeError,
2941 "can only join an iterable of bytes "
2942 "(item %ld has type '%.100s')",
2943 /* XXX %ld isn't right on Win64 */
2944 (long)i, Py_TYPE(obj)->tp_name);
2945 goto error;
2946 }
2947 if (i > 0)
2948 totalsize += mysize;
2949 totalsize += Py_SIZE(obj);
2950 if (totalsize < 0) {
2951 PyErr_NoMemory();
2952 goto error;
2953 }
2954 }
2955
2956 /* Allocate the result, and copy the bytes */
2957 result = PyByteArray_FromStringAndSize(NULL, totalsize);
2958 if (result == NULL)
2959 goto error;
2960 dest = PyByteArray_AS_STRING(result);
2961 for (i = 0; i < n; i++) {
2962 PyObject *obj = items[i];
2963 Py_ssize_t size = Py_SIZE(obj);
2964 char *buf;
2965 if (PyByteArray_Check(obj))
2966 buf = PyByteArray_AS_STRING(obj);
2967 else
2968 buf = PyBytes_AS_STRING(obj);
2969 if (i) {
2970 memcpy(dest, self->ob_bytes, mysize);
2971 dest += mysize;
2972 }
2973 memcpy(dest, buf, size);
2974 dest += size;
2975 }
2976
2977 /* Done */
2978 Py_DECREF(seq);
2979 return result;
2980
2981 /* Error handling */
2982 error:
2983 Py_DECREF(seq);
2984 return NULL;
2985}
2986
2987PyDoc_STRVAR(fromhex_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002988"bytearray.fromhex(string) -> bytearray (static method)\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002989\n\
2990Create a bytearray object from a string of hexadecimal numbers.\n\
2991Spaces between two numbers are accepted.\n\
2992Example: bytearray.fromhex('B9 01EF') -> bytearray(b'\\xb9\\x01\\xef').");
2993
2994static int
2995hex_digit_to_int(Py_UNICODE c)
2996{
2997 if (c >= 128)
2998 return -1;
Eric Smith6dc46f52009-04-27 20:39:49 +00002999 if (Py_ISDIGIT(c))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003000 return c - '0';
3001 else {
Eric Smith6dc46f52009-04-27 20:39:49 +00003002 if (Py_ISUPPER(c))
3003 c = Py_TOLOWER(c);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003004 if (c >= 'a' && c <= 'f')
3005 return c - 'a' + 10;
3006 }
3007 return -1;
3008}
3009
3010static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003011bytearray_fromhex(PyObject *cls, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003012{
3013 PyObject *newbytes, *hexobj;
3014 char *buf;
3015 Py_UNICODE *hex;
3016 Py_ssize_t hexlen, byteslen, i, j;
3017 int top, bot;
3018
3019 if (!PyArg_ParseTuple(args, "U:fromhex", &hexobj))
3020 return NULL;
3021 assert(PyUnicode_Check(hexobj));
3022 hexlen = PyUnicode_GET_SIZE(hexobj);
3023 hex = PyUnicode_AS_UNICODE(hexobj);
3024 byteslen = hexlen/2; /* This overestimates if there are spaces */
3025 newbytes = PyByteArray_FromStringAndSize(NULL, byteslen);
3026 if (!newbytes)
3027 return NULL;
3028 buf = PyByteArray_AS_STRING(newbytes);
3029 for (i = j = 0; i < hexlen; i += 2) {
3030 /* skip over spaces in the input */
3031 while (hex[i] == ' ')
3032 i++;
3033 if (i >= hexlen)
3034 break;
3035 top = hex_digit_to_int(hex[i]);
3036 bot = hex_digit_to_int(hex[i+1]);
3037 if (top == -1 || bot == -1) {
3038 PyErr_Format(PyExc_ValueError,
3039 "non-hexadecimal number found in "
3040 "fromhex() arg at position %zd", i);
3041 goto error;
3042 }
3043 buf[j++] = (top << 4) + bot;
3044 }
3045 if (PyByteArray_Resize(newbytes, j) < 0)
3046 goto error;
3047 return newbytes;
3048
3049 error:
3050 Py_DECREF(newbytes);
3051 return NULL;
3052}
3053
3054PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
3055
3056static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003057bytearray_reduce(PyByteArrayObject *self)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003058{
3059 PyObject *latin1, *dict;
3060 if (self->ob_bytes)
3061 latin1 = PyUnicode_DecodeLatin1(self->ob_bytes,
3062 Py_SIZE(self), NULL);
3063 else
3064 latin1 = PyUnicode_FromString("");
3065
3066 dict = PyObject_GetAttrString((PyObject *)self, "__dict__");
3067 if (dict == NULL) {
3068 PyErr_Clear();
3069 dict = Py_None;
3070 Py_INCREF(dict);
3071 }
3072
3073 return Py_BuildValue("(O(Ns)N)", Py_TYPE(self), latin1, "latin-1", dict);
3074}
3075
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003076PyDoc_STRVAR(sizeof_doc,
3077"B.__sizeof__() -> int\n\
3078 \n\
3079Returns the size of B in memory, in bytes");
3080static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003081bytearray_sizeof(PyByteArrayObject *self)
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003082{
Benjamin Petersonef3e4c22009-04-11 19:48:14 +00003083 Py_ssize_t res;
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003084
Benjamin Petersonef3e4c22009-04-11 19:48:14 +00003085 res = sizeof(PyByteArrayObject) + self->ob_alloc * sizeof(char);
3086 return PyLong_FromSsize_t(res);
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003087}
3088
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003089static PySequenceMethods bytearray_as_sequence = {
3090 (lenfunc)bytearray_length, /* sq_length */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003091 (binaryfunc)PyByteArray_Concat, /* sq_concat */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003092 (ssizeargfunc)bytearray_repeat, /* sq_repeat */
3093 (ssizeargfunc)bytearray_getitem, /* sq_item */
3094 0, /* sq_slice */
3095 (ssizeobjargproc)bytearray_setitem, /* sq_ass_item */
3096 0, /* sq_ass_slice */
3097 (objobjproc)bytearray_contains, /* sq_contains */
3098 (binaryfunc)bytearray_iconcat, /* sq_inplace_concat */
3099 (ssizeargfunc)bytearray_irepeat, /* sq_inplace_repeat */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003100};
3101
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003102static PyMappingMethods bytearray_as_mapping = {
3103 (lenfunc)bytearray_length,
3104 (binaryfunc)bytearray_subscript,
3105 (objobjargproc)bytearray_ass_subscript,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003106};
3107
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003108static PyBufferProcs bytearray_as_buffer = {
3109 (getbufferproc)bytearray_getbuffer,
3110 (releasebufferproc)bytearray_releasebuffer,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003111};
3112
3113static PyMethodDef
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003114bytearray_methods[] = {
3115 {"__alloc__", (PyCFunction)bytearray_alloc, METH_NOARGS, alloc_doc},
3116 {"__reduce__", (PyCFunction)bytearray_reduce, METH_NOARGS, reduce_doc},
3117 {"__sizeof__", (PyCFunction)bytearray_sizeof, METH_NOARGS, sizeof_doc},
3118 {"append", (PyCFunction)bytearray_append, METH_O, append__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003119 {"capitalize", (PyCFunction)stringlib_capitalize, METH_NOARGS,
3120 _Py_capitalize__doc__},
3121 {"center", (PyCFunction)stringlib_center, METH_VARARGS, center__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003122 {"count", (PyCFunction)bytearray_count, METH_VARARGS, count__doc__},
3123 {"decode", (PyCFunction)bytearray_decode, METH_VARARGS, decode_doc},
3124 {"endswith", (PyCFunction)bytearray_endswith, METH_VARARGS, endswith__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003125 {"expandtabs", (PyCFunction)stringlib_expandtabs, METH_VARARGS,
3126 expandtabs__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003127 {"extend", (PyCFunction)bytearray_extend, METH_O, extend__doc__},
3128 {"find", (PyCFunction)bytearray_find, METH_VARARGS, find__doc__},
3129 {"fromhex", (PyCFunction)bytearray_fromhex, METH_VARARGS|METH_CLASS,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003130 fromhex_doc},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003131 {"index", (PyCFunction)bytearray_index, METH_VARARGS, index__doc__},
3132 {"insert", (PyCFunction)bytearray_insert, METH_VARARGS, insert__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003133 {"isalnum", (PyCFunction)stringlib_isalnum, METH_NOARGS,
3134 _Py_isalnum__doc__},
3135 {"isalpha", (PyCFunction)stringlib_isalpha, METH_NOARGS,
3136 _Py_isalpha__doc__},
3137 {"isdigit", (PyCFunction)stringlib_isdigit, METH_NOARGS,
3138 _Py_isdigit__doc__},
3139 {"islower", (PyCFunction)stringlib_islower, METH_NOARGS,
3140 _Py_islower__doc__},
3141 {"isspace", (PyCFunction)stringlib_isspace, METH_NOARGS,
3142 _Py_isspace__doc__},
3143 {"istitle", (PyCFunction)stringlib_istitle, METH_NOARGS,
3144 _Py_istitle__doc__},
3145 {"isupper", (PyCFunction)stringlib_isupper, METH_NOARGS,
3146 _Py_isupper__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003147 {"join", (PyCFunction)bytearray_join, METH_O, join_doc},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003148 {"ljust", (PyCFunction)stringlib_ljust, METH_VARARGS, ljust__doc__},
3149 {"lower", (PyCFunction)stringlib_lower, METH_NOARGS, _Py_lower__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003150 {"lstrip", (PyCFunction)bytearray_lstrip, METH_VARARGS, lstrip__doc__},
3151 {"maketrans", (PyCFunction)bytearray_maketrans, METH_VARARGS|METH_STATIC,
Georg Brandlabc38772009-04-12 15:51:51 +00003152 _Py_maketrans__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003153 {"partition", (PyCFunction)bytearray_partition, METH_O, partition__doc__},
3154 {"pop", (PyCFunction)bytearray_pop, METH_VARARGS, pop__doc__},
3155 {"remove", (PyCFunction)bytearray_remove, METH_O, remove__doc__},
3156 {"replace", (PyCFunction)bytearray_replace, METH_VARARGS, replace__doc__},
3157 {"reverse", (PyCFunction)bytearray_reverse, METH_NOARGS, reverse__doc__},
3158 {"rfind", (PyCFunction)bytearray_rfind, METH_VARARGS, rfind__doc__},
3159 {"rindex", (PyCFunction)bytearray_rindex, METH_VARARGS, rindex__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003160 {"rjust", (PyCFunction)stringlib_rjust, METH_VARARGS, rjust__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003161 {"rpartition", (PyCFunction)bytearray_rpartition, METH_O, rpartition__doc__},
3162 {"rsplit", (PyCFunction)bytearray_rsplit, METH_VARARGS, rsplit__doc__},
3163 {"rstrip", (PyCFunction)bytearray_rstrip, METH_VARARGS, rstrip__doc__},
3164 {"split", (PyCFunction)bytearray_split, METH_VARARGS, split__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003165 {"splitlines", (PyCFunction)stringlib_splitlines, METH_VARARGS,
3166 splitlines__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003167 {"startswith", (PyCFunction)bytearray_startswith, METH_VARARGS ,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003168 startswith__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003169 {"strip", (PyCFunction)bytearray_strip, METH_VARARGS, strip__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003170 {"swapcase", (PyCFunction)stringlib_swapcase, METH_NOARGS,
3171 _Py_swapcase__doc__},
3172 {"title", (PyCFunction)stringlib_title, METH_NOARGS, _Py_title__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003173 {"translate", (PyCFunction)bytearray_translate, METH_VARARGS,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003174 translate__doc__},
3175 {"upper", (PyCFunction)stringlib_upper, METH_NOARGS, _Py_upper__doc__},
3176 {"zfill", (PyCFunction)stringlib_zfill, METH_VARARGS, zfill__doc__},
3177 {NULL}
3178};
3179
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003180PyDoc_STRVAR(bytearray_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00003181"bytearray(iterable_of_ints) -> bytearray\n\
3182bytearray(string, encoding[, errors]) -> bytearray\n\
3183bytearray(bytes_or_bytearray) -> mutable copy of bytes_or_bytearray\n\
3184bytearray(memory_view) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003185\n\
3186Construct an mutable bytearray object from:\n\
3187 - an iterable yielding integers in range(256)\n\
3188 - a text string encoded using the specified encoding\n\
3189 - a bytes or a bytearray object\n\
3190 - any object implementing the buffer API.\n\
3191\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00003192bytearray(int) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003193\n\
3194Construct a zero-initialized bytearray of the given length.");
3195
3196
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003197static PyObject *bytearray_iter(PyObject *seq);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003198
3199PyTypeObject PyByteArray_Type = {
3200 PyVarObject_HEAD_INIT(&PyType_Type, 0)
3201 "bytearray",
3202 sizeof(PyByteArrayObject),
3203 0,
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003204 (destructor)bytearray_dealloc, /* tp_dealloc */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003205 0, /* tp_print */
3206 0, /* tp_getattr */
3207 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003208 0, /* tp_reserved */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003209 (reprfunc)bytearray_repr, /* tp_repr */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003210 0, /* tp_as_number */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003211 &bytearray_as_sequence, /* tp_as_sequence */
3212 &bytearray_as_mapping, /* tp_as_mapping */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003213 0, /* tp_hash */
3214 0, /* tp_call */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003215 bytearray_str, /* tp_str */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003216 PyObject_GenericGetAttr, /* tp_getattro */
3217 0, /* tp_setattro */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003218 &bytearray_as_buffer, /* tp_as_buffer */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003219 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003220 bytearray_doc, /* tp_doc */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003221 0, /* tp_traverse */
3222 0, /* tp_clear */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003223 (richcmpfunc)bytearray_richcompare, /* tp_richcompare */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003224 0, /* tp_weaklistoffset */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003225 bytearray_iter, /* tp_iter */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003226 0, /* tp_iternext */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003227 bytearray_methods, /* tp_methods */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003228 0, /* tp_members */
3229 0, /* tp_getset */
3230 0, /* tp_base */
3231 0, /* tp_dict */
3232 0, /* tp_descr_get */
3233 0, /* tp_descr_set */
3234 0, /* tp_dictoffset */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003235 (initproc)bytearray_init, /* tp_init */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003236 PyType_GenericAlloc, /* tp_alloc */
3237 PyType_GenericNew, /* tp_new */
3238 PyObject_Del, /* tp_free */
3239};
3240
3241/*********************** Bytes Iterator ****************************/
3242
3243typedef struct {
3244 PyObject_HEAD
3245 Py_ssize_t it_index;
3246 PyByteArrayObject *it_seq; /* Set to NULL when iterator is exhausted */
3247} bytesiterobject;
3248
3249static void
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003250bytearrayiter_dealloc(bytesiterobject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003251{
3252 _PyObject_GC_UNTRACK(it);
3253 Py_XDECREF(it->it_seq);
3254 PyObject_GC_Del(it);
3255}
3256
3257static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003258bytearrayiter_traverse(bytesiterobject *it, visitproc visit, void *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003259{
3260 Py_VISIT(it->it_seq);
3261 return 0;
3262}
3263
3264static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003265bytearrayiter_next(bytesiterobject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003266{
3267 PyByteArrayObject *seq;
3268 PyObject *item;
3269
3270 assert(it != NULL);
3271 seq = it->it_seq;
3272 if (seq == NULL)
3273 return NULL;
3274 assert(PyByteArray_Check(seq));
3275
3276 if (it->it_index < PyByteArray_GET_SIZE(seq)) {
3277 item = PyLong_FromLong(
3278 (unsigned char)seq->ob_bytes[it->it_index]);
3279 if (item != NULL)
3280 ++it->it_index;
3281 return item;
3282 }
3283
3284 Py_DECREF(seq);
3285 it->it_seq = NULL;
3286 return NULL;
3287}
3288
3289static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003290bytesarrayiter_length_hint(bytesiterobject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003291{
3292 Py_ssize_t len = 0;
3293 if (it->it_seq)
3294 len = PyByteArray_GET_SIZE(it->it_seq) - it->it_index;
3295 return PyLong_FromSsize_t(len);
3296}
3297
3298PyDoc_STRVAR(length_hint_doc,
3299 "Private method returning an estimate of len(list(it)).");
3300
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003301static PyMethodDef bytearrayiter_methods[] = {
3302 {"__length_hint__", (PyCFunction)bytesarrayiter_length_hint, METH_NOARGS,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003303 length_hint_doc},
3304 {NULL, NULL} /* sentinel */
3305};
3306
3307PyTypeObject PyByteArrayIter_Type = {
3308 PyVarObject_HEAD_INIT(&PyType_Type, 0)
3309 "bytearray_iterator", /* tp_name */
3310 sizeof(bytesiterobject), /* tp_basicsize */
3311 0, /* tp_itemsize */
3312 /* methods */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003313 (destructor)bytearrayiter_dealloc, /* tp_dealloc */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003314 0, /* tp_print */
3315 0, /* tp_getattr */
3316 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003317 0, /* tp_reserved */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003318 0, /* tp_repr */
3319 0, /* tp_as_number */
3320 0, /* tp_as_sequence */
3321 0, /* tp_as_mapping */
3322 0, /* tp_hash */
3323 0, /* tp_call */
3324 0, /* tp_str */
3325 PyObject_GenericGetAttr, /* tp_getattro */
3326 0, /* tp_setattro */
3327 0, /* tp_as_buffer */
3328 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
3329 0, /* tp_doc */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003330 (traverseproc)bytearrayiter_traverse, /* tp_traverse */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003331 0, /* tp_clear */
3332 0, /* tp_richcompare */
3333 0, /* tp_weaklistoffset */
3334 PyObject_SelfIter, /* tp_iter */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003335 (iternextfunc)bytearrayiter_next, /* tp_iternext */
3336 bytearrayiter_methods, /* tp_methods */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003337 0,
3338};
3339
3340static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003341bytearray_iter(PyObject *seq)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003342{
3343 bytesiterobject *it;
3344
3345 if (!PyByteArray_Check(seq)) {
3346 PyErr_BadInternalCall();
3347 return NULL;
3348 }
3349 it = PyObject_GC_New(bytesiterobject, &PyByteArrayIter_Type);
3350 if (it == NULL)
3351 return NULL;
3352 it->it_index = 0;
3353 Py_INCREF(seq);
3354 it->it_seq = (PyByteArrayObject *)seq;
3355 _PyObject_GC_TRACK(it);
3356 return (PyObject *)it;
3357}