blob: 670f5899958e70781e0ab106579ed4ba6f8eb82e [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)) {
598 /* Make a copy an call this function recursively */
599 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{
939 if (Py_BytesWarningFlag) {
940 if (PyErr_WarnEx(PyExc_BytesWarning,
941 "str() on a bytearray instance", 1))
942 return NULL;
943 }
Benjamin Peterson153c70f2009-04-18 15:42:12 +0000944 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
1085 if (!PyArg_ParseTuple(args, "O|O&O&:find/rfind/index/rindex", &subobj,
1086 _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1087 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
1136 if (!PyArg_ParseTuple(args, "O|O&O&:count", &sub_obj,
1137 _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1138 return NULL;
1139
1140 if (_getbuffer(sub_obj, &vsub) < 0)
1141 return NULL;
1142
1143 _adjust_indices(&start, &end, PyByteArray_GET_SIZE(self));
1144
1145 count_obj = PyLong_FromSsize_t(
1146 stringlib_count(str + start, end - start, vsub.buf, vsub.len)
1147 );
Martin v. Löwis423be952008-08-13 15:53:07 +00001148 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001149 return count_obj;
1150}
1151
1152
1153PyDoc_STRVAR(index__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001154"B.index(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001155\n\
1156Like B.find() but raise ValueError when the subsection is not found.");
1157
1158static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001159bytearray_index(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001160{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001161 Py_ssize_t result = bytearray_find_internal(self, args, +1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001162 if (result == -2)
1163 return NULL;
1164 if (result == -1) {
1165 PyErr_SetString(PyExc_ValueError,
1166 "subsection not found");
1167 return NULL;
1168 }
1169 return PyLong_FromSsize_t(result);
1170}
1171
1172
1173PyDoc_STRVAR(rfind__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001174"B.rfind(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001175\n\
1176Return the highest index in B where subsection sub is found,\n\
1177such that sub is contained within s[start,end]. Optional\n\
1178arguments start and end are interpreted as in slice notation.\n\
1179\n\
1180Return -1 on failure.");
1181
1182static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001183bytearray_rfind(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001184{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001185 Py_ssize_t result = bytearray_find_internal(self, args, -1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001186 if (result == -2)
1187 return NULL;
1188 return PyLong_FromSsize_t(result);
1189}
1190
1191
1192PyDoc_STRVAR(rindex__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001193"B.rindex(sub[, start[, end]]) -> int\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001194\n\
1195Like B.rfind() but raise ValueError when the subsection is not found.");
1196
1197static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001198bytearray_rindex(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001199{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001200 Py_ssize_t result = bytearray_find_internal(self, args, -1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001201 if (result == -2)
1202 return NULL;
1203 if (result == -1) {
1204 PyErr_SetString(PyExc_ValueError,
1205 "subsection not found");
1206 return NULL;
1207 }
1208 return PyLong_FromSsize_t(result);
1209}
1210
1211
1212static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001213bytearray_contains(PyObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001214{
1215 Py_ssize_t ival = PyNumber_AsSsize_t(arg, PyExc_ValueError);
1216 if (ival == -1 && PyErr_Occurred()) {
1217 Py_buffer varg;
Antoine Pitroubc760d92010-08-15 17:46:50 +00001218 Py_ssize_t pos;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001219 PyErr_Clear();
1220 if (_getbuffer(arg, &varg) < 0)
1221 return -1;
1222 pos = stringlib_find(PyByteArray_AS_STRING(self), Py_SIZE(self),
1223 varg.buf, varg.len, 0);
Martin v. Löwis423be952008-08-13 15:53:07 +00001224 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001225 return pos >= 0;
1226 }
1227 if (ival < 0 || ival >= 256) {
1228 PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
1229 return -1;
1230 }
1231
Antoine Pitroubc760d92010-08-15 17:46:50 +00001232 return memchr(PyByteArray_AS_STRING(self), (int) ival, Py_SIZE(self)) != NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001233}
1234
1235
1236/* Matches the end (direction >= 0) or start (direction < 0) of self
1237 * against substr, using the start and end arguments. Returns
1238 * -1 on error, 0 if not found and 1 if found.
1239 */
1240Py_LOCAL(int)
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001241_bytearray_tailmatch(PyByteArrayObject *self, PyObject *substr, Py_ssize_t start,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001242 Py_ssize_t end, int direction)
1243{
1244 Py_ssize_t len = PyByteArray_GET_SIZE(self);
1245 const char* str;
1246 Py_buffer vsubstr;
1247 int rv = 0;
1248
1249 str = PyByteArray_AS_STRING(self);
1250
1251 if (_getbuffer(substr, &vsubstr) < 0)
1252 return -1;
1253
1254 _adjust_indices(&start, &end, len);
1255
1256 if (direction < 0) {
1257 /* startswith */
1258 if (start+vsubstr.len > len) {
1259 goto done;
1260 }
1261 } else {
1262 /* endswith */
1263 if (end-start < vsubstr.len || start > len) {
1264 goto done;
1265 }
1266
1267 if (end-vsubstr.len > start)
1268 start = end - vsubstr.len;
1269 }
1270 if (end-start >= vsubstr.len)
1271 rv = ! memcmp(str+start, vsubstr.buf, vsubstr.len);
1272
1273done:
Martin v. Löwis423be952008-08-13 15:53:07 +00001274 PyBuffer_Release(&vsubstr);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001275 return rv;
1276}
1277
1278
1279PyDoc_STRVAR(startswith__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001280"B.startswith(prefix[, start[, end]]) -> bool\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001281\n\
1282Return True if B starts with the specified prefix, False otherwise.\n\
1283With optional start, test B beginning at that position.\n\
1284With optional end, stop comparing B at that position.\n\
1285prefix can also be a tuple of strings to try.");
1286
1287static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001288bytearray_startswith(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001289{
1290 Py_ssize_t start = 0;
1291 Py_ssize_t end = PY_SSIZE_T_MAX;
1292 PyObject *subobj;
1293 int result;
1294
1295 if (!PyArg_ParseTuple(args, "O|O&O&:startswith", &subobj,
1296 _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1297 return NULL;
1298 if (PyTuple_Check(subobj)) {
1299 Py_ssize_t i;
1300 for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001301 result = _bytearray_tailmatch(self,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001302 PyTuple_GET_ITEM(subobj, i),
1303 start, end, -1);
1304 if (result == -1)
1305 return NULL;
1306 else if (result) {
1307 Py_RETURN_TRUE;
1308 }
1309 }
1310 Py_RETURN_FALSE;
1311 }
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001312 result = _bytearray_tailmatch(self, subobj, start, end, -1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001313 if (result == -1)
1314 return NULL;
1315 else
1316 return PyBool_FromLong(result);
1317}
1318
1319PyDoc_STRVAR(endswith__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00001320"B.endswith(suffix[, start[, end]]) -> bool\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001321\n\
1322Return True if B ends with the specified suffix, False otherwise.\n\
1323With optional start, test B beginning at that position.\n\
1324With optional end, stop comparing B at that position.\n\
1325suffix can also be a tuple of strings to try.");
1326
1327static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001328bytearray_endswith(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001329{
1330 Py_ssize_t start = 0;
1331 Py_ssize_t end = PY_SSIZE_T_MAX;
1332 PyObject *subobj;
1333 int result;
1334
1335 if (!PyArg_ParseTuple(args, "O|O&O&:endswith", &subobj,
1336 _PyEval_SliceIndex, &start, _PyEval_SliceIndex, &end))
1337 return NULL;
1338 if (PyTuple_Check(subobj)) {
1339 Py_ssize_t i;
1340 for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001341 result = _bytearray_tailmatch(self,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001342 PyTuple_GET_ITEM(subobj, i),
1343 start, end, +1);
1344 if (result == -1)
1345 return NULL;
1346 else if (result) {
1347 Py_RETURN_TRUE;
1348 }
1349 }
1350 Py_RETURN_FALSE;
1351 }
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001352 result = _bytearray_tailmatch(self, subobj, start, end, +1);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001353 if (result == -1)
1354 return NULL;
1355 else
1356 return PyBool_FromLong(result);
1357}
1358
1359
1360PyDoc_STRVAR(translate__doc__,
1361"B.translate(table[, deletechars]) -> bytearray\n\
1362\n\
1363Return a copy of B, where all characters occurring in the\n\
1364optional argument deletechars are removed, and the remaining\n\
1365characters have been mapped through the given translation\n\
1366table, which must be a bytes object of length 256.");
1367
1368static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001369bytearray_translate(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001370{
1371 register char *input, *output;
1372 register const char *table;
Benjamin Peterson0f3641c2008-11-19 22:05:52 +00001373 register Py_ssize_t i, c;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001374 PyObject *input_obj = (PyObject*)self;
1375 const char *output_start;
1376 Py_ssize_t inlen;
Georg Brandlccc47b62008-12-28 11:44:14 +00001377 PyObject *result = NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001378 int trans_table[256];
Georg Brandlccc47b62008-12-28 11:44:14 +00001379 PyObject *tableobj = NULL, *delobj = NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001380 Py_buffer vtable, vdel;
1381
1382 if (!PyArg_UnpackTuple(args, "translate", 1, 2,
1383 &tableobj, &delobj))
1384 return NULL;
1385
Georg Brandlccc47b62008-12-28 11:44:14 +00001386 if (tableobj == Py_None) {
1387 table = NULL;
1388 tableobj = NULL;
1389 } else if (_getbuffer(tableobj, &vtable) < 0) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001390 return NULL;
Georg Brandlccc47b62008-12-28 11:44:14 +00001391 } else {
1392 if (vtable.len != 256) {
1393 PyErr_SetString(PyExc_ValueError,
1394 "translation table must be 256 characters long");
Georg Brandl069bcc32009-07-22 12:06:11 +00001395 PyBuffer_Release(&vtable);
1396 return NULL;
Georg Brandlccc47b62008-12-28 11:44:14 +00001397 }
1398 table = (const char*)vtable.buf;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001399 }
1400
1401 if (delobj != NULL) {
1402 if (_getbuffer(delobj, &vdel) < 0) {
Georg Brandl069bcc32009-07-22 12:06:11 +00001403 if (tableobj != NULL)
1404 PyBuffer_Release(&vtable);
1405 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001406 }
1407 }
1408 else {
1409 vdel.buf = NULL;
1410 vdel.len = 0;
1411 }
1412
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001413 inlen = PyByteArray_GET_SIZE(input_obj);
1414 result = PyByteArray_FromStringAndSize((char *)NULL, inlen);
1415 if (result == NULL)
1416 goto done;
1417 output_start = output = PyByteArray_AsString(result);
1418 input = PyByteArray_AS_STRING(input_obj);
1419
Georg Brandlccc47b62008-12-28 11:44:14 +00001420 if (vdel.len == 0 && table != NULL) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001421 /* If no deletions are required, use faster code */
1422 for (i = inlen; --i >= 0; ) {
1423 c = Py_CHARMASK(*input++);
Benjamin Peterson0f3641c2008-11-19 22:05:52 +00001424 *output++ = table[c];
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001425 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001426 goto done;
1427 }
Georg Brandlccc47b62008-12-28 11:44:14 +00001428
1429 if (table == NULL) {
1430 for (i = 0; i < 256; i++)
1431 trans_table[i] = Py_CHARMASK(i);
1432 } else {
1433 for (i = 0; i < 256; i++)
1434 trans_table[i] = Py_CHARMASK(table[i]);
1435 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001436
1437 for (i = 0; i < vdel.len; i++)
1438 trans_table[(int) Py_CHARMASK( ((unsigned char*)vdel.buf)[i] )] = -1;
1439
1440 for (i = inlen; --i >= 0; ) {
1441 c = Py_CHARMASK(*input++);
1442 if (trans_table[c] != -1)
1443 if (Py_CHARMASK(*output++ = (char)trans_table[c]) == c)
1444 continue;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001445 }
1446 /* Fix the size of the resulting string */
1447 if (inlen > 0)
1448 PyByteArray_Resize(result, output - output_start);
1449
1450done:
Georg Brandlccc47b62008-12-28 11:44:14 +00001451 if (tableobj != NULL)
1452 PyBuffer_Release(&vtable);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001453 if (delobj != NULL)
Martin v. Löwis423be952008-08-13 15:53:07 +00001454 PyBuffer_Release(&vdel);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001455 return result;
1456}
1457
1458
Georg Brandlabc38772009-04-12 15:51:51 +00001459static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00001460bytearray_maketrans(PyObject *null, PyObject *args)
Georg Brandlabc38772009-04-12 15:51:51 +00001461{
1462 return _Py_bytes_maketrans(args);
1463}
1464
1465
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001466#define FORWARD 1
1467#define REVERSE -1
1468
1469/* find and count characters and substrings */
1470
1471#define findchar(target, target_len, c) \
1472 ((char *)memchr((const void *)(target), c, target_len))
1473
1474/* Don't call if length < 2 */
1475#define Py_STRING_MATCH(target, offset, pattern, length) \
1476 (target[offset] == pattern[0] && \
1477 target[offset+length-1] == pattern[length-1] && \
1478 !memcmp(target+offset+1, pattern+1, length-2) )
1479
1480
Benjamin Peterson0f3641c2008-11-19 22:05:52 +00001481/* Bytes ops must return a string, create a copy */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001482Py_LOCAL(PyByteArrayObject *)
1483return_self(PyByteArrayObject *self)
1484{
Georg Brandl1e7217d2008-05-30 12:02:38 +00001485 /* always return a new bytearray */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001486 return (PyByteArrayObject *)PyByteArray_FromStringAndSize(
1487 PyByteArray_AS_STRING(self),
1488 PyByteArray_GET_SIZE(self));
1489}
1490
1491Py_LOCAL_INLINE(Py_ssize_t)
1492countchar(const char *target, Py_ssize_t target_len, char c, Py_ssize_t maxcount)
1493{
1494 Py_ssize_t count=0;
1495 const char *start=target;
1496 const char *end=target+target_len;
1497
1498 while ( (start=findchar(start, end-start, c)) != NULL ) {
1499 count++;
1500 if (count >= maxcount)
1501 break;
1502 start += 1;
1503 }
1504 return count;
1505}
1506
1507Py_LOCAL(Py_ssize_t)
1508findstring(const char *target, Py_ssize_t target_len,
1509 const char *pattern, Py_ssize_t pattern_len,
1510 Py_ssize_t start,
1511 Py_ssize_t end,
1512 int direction)
1513{
1514 if (start < 0) {
1515 start += target_len;
1516 if (start < 0)
1517 start = 0;
1518 }
1519 if (end > target_len) {
1520 end = target_len;
1521 } else if (end < 0) {
1522 end += target_len;
1523 if (end < 0)
1524 end = 0;
1525 }
1526
1527 /* zero-length substrings always match at the first attempt */
1528 if (pattern_len == 0)
1529 return (direction > 0) ? start : end;
1530
1531 end -= pattern_len;
1532
1533 if (direction < 0) {
1534 for (; end >= start; end--)
1535 if (Py_STRING_MATCH(target, end, pattern, pattern_len))
1536 return end;
1537 } else {
1538 for (; start <= end; start++)
1539 if (Py_STRING_MATCH(target, start, pattern, pattern_len))
1540 return start;
1541 }
1542 return -1;
1543}
1544
1545Py_LOCAL_INLINE(Py_ssize_t)
1546countstring(const char *target, Py_ssize_t target_len,
1547 const char *pattern, Py_ssize_t pattern_len,
1548 Py_ssize_t start,
1549 Py_ssize_t end,
1550 int direction, Py_ssize_t maxcount)
1551{
1552 Py_ssize_t count=0;
1553
1554 if (start < 0) {
1555 start += target_len;
1556 if (start < 0)
1557 start = 0;
1558 }
1559 if (end > target_len) {
1560 end = target_len;
1561 } else if (end < 0) {
1562 end += target_len;
1563 if (end < 0)
1564 end = 0;
1565 }
1566
1567 /* zero-length substrings match everywhere */
1568 if (pattern_len == 0 || maxcount == 0) {
1569 if (target_len+1 < maxcount)
1570 return target_len+1;
1571 return maxcount;
1572 }
1573
1574 end -= pattern_len;
1575 if (direction < 0) {
1576 for (; (end >= start); end--)
1577 if (Py_STRING_MATCH(target, end, pattern, pattern_len)) {
1578 count++;
1579 if (--maxcount <= 0) break;
1580 end -= pattern_len-1;
1581 }
1582 } else {
1583 for (; (start <= end); start++)
1584 if (Py_STRING_MATCH(target, start, pattern, pattern_len)) {
1585 count++;
1586 if (--maxcount <= 0)
1587 break;
1588 start += pattern_len-1;
1589 }
1590 }
1591 return count;
1592}
1593
1594
1595/* Algorithms for different cases of string replacement */
1596
1597/* len(self)>=1, from="", len(to)>=1, maxcount>=1 */
1598Py_LOCAL(PyByteArrayObject *)
1599replace_interleave(PyByteArrayObject *self,
1600 const char *to_s, Py_ssize_t to_len,
1601 Py_ssize_t maxcount)
1602{
1603 char *self_s, *result_s;
1604 Py_ssize_t self_len, result_len;
1605 Py_ssize_t count, i, product;
1606 PyByteArrayObject *result;
1607
1608 self_len = PyByteArray_GET_SIZE(self);
1609
1610 /* 1 at the end plus 1 after every character */
1611 count = self_len+1;
1612 if (maxcount < count)
1613 count = maxcount;
1614
1615 /* Check for overflow */
1616 /* result_len = count * to_len + self_len; */
1617 product = count * to_len;
1618 if (product / to_len != count) {
1619 PyErr_SetString(PyExc_OverflowError,
1620 "replace string is too long");
1621 return NULL;
1622 }
1623 result_len = product + self_len;
1624 if (result_len < 0) {
1625 PyErr_SetString(PyExc_OverflowError,
1626 "replace string is too long");
1627 return NULL;
1628 }
1629
1630 if (! (result = (PyByteArrayObject *)
1631 PyByteArray_FromStringAndSize(NULL, result_len)) )
1632 return NULL;
1633
1634 self_s = PyByteArray_AS_STRING(self);
1635 result_s = PyByteArray_AS_STRING(result);
1636
1637 /* TODO: special case single character, which doesn't need memcpy */
1638
1639 /* Lay the first one down (guaranteed this will occur) */
1640 Py_MEMCPY(result_s, to_s, to_len);
1641 result_s += to_len;
1642 count -= 1;
1643
1644 for (i=0; i<count; i++) {
1645 *result_s++ = *self_s++;
1646 Py_MEMCPY(result_s, to_s, to_len);
1647 result_s += to_len;
1648 }
1649
1650 /* Copy the rest of the original string */
1651 Py_MEMCPY(result_s, self_s, self_len-i);
1652
1653 return result;
1654}
1655
1656/* Special case for deleting a single character */
1657/* len(self)>=1, len(from)==1, to="", maxcount>=1 */
1658Py_LOCAL(PyByteArrayObject *)
1659replace_delete_single_character(PyByteArrayObject *self,
1660 char from_c, Py_ssize_t maxcount)
1661{
1662 char *self_s, *result_s;
1663 char *start, *next, *end;
1664 Py_ssize_t self_len, result_len;
1665 Py_ssize_t count;
1666 PyByteArrayObject *result;
1667
1668 self_len = PyByteArray_GET_SIZE(self);
1669 self_s = PyByteArray_AS_STRING(self);
1670
1671 count = countchar(self_s, self_len, from_c, maxcount);
1672 if (count == 0) {
1673 return return_self(self);
1674 }
1675
1676 result_len = self_len - count; /* from_len == 1 */
1677 assert(result_len>=0);
1678
1679 if ( (result = (PyByteArrayObject *)
1680 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL)
1681 return NULL;
1682 result_s = PyByteArray_AS_STRING(result);
1683
1684 start = self_s;
1685 end = self_s + self_len;
1686 while (count-- > 0) {
1687 next = findchar(start, end-start, from_c);
1688 if (next == NULL)
1689 break;
1690 Py_MEMCPY(result_s, start, next-start);
1691 result_s += (next-start);
1692 start = next+1;
1693 }
1694 Py_MEMCPY(result_s, start, end-start);
1695
1696 return result;
1697}
1698
1699/* len(self)>=1, len(from)>=2, to="", maxcount>=1 */
1700
1701Py_LOCAL(PyByteArrayObject *)
1702replace_delete_substring(PyByteArrayObject *self,
1703 const char *from_s, Py_ssize_t from_len,
1704 Py_ssize_t maxcount)
1705{
1706 char *self_s, *result_s;
1707 char *start, *next, *end;
1708 Py_ssize_t self_len, result_len;
1709 Py_ssize_t count, offset;
1710 PyByteArrayObject *result;
1711
1712 self_len = PyByteArray_GET_SIZE(self);
1713 self_s = PyByteArray_AS_STRING(self);
1714
1715 count = countstring(self_s, self_len,
1716 from_s, from_len,
1717 0, self_len, 1,
1718 maxcount);
1719
1720 if (count == 0) {
1721 /* no matches */
1722 return return_self(self);
1723 }
1724
1725 result_len = self_len - (count * from_len);
1726 assert (result_len>=0);
1727
1728 if ( (result = (PyByteArrayObject *)
1729 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL )
1730 return NULL;
1731
1732 result_s = PyByteArray_AS_STRING(result);
1733
1734 start = self_s;
1735 end = self_s + self_len;
1736 while (count-- > 0) {
1737 offset = findstring(start, end-start,
1738 from_s, from_len,
1739 0, end-start, FORWARD);
1740 if (offset == -1)
1741 break;
1742 next = start + offset;
1743
1744 Py_MEMCPY(result_s, start, next-start);
1745
1746 result_s += (next-start);
1747 start = next+from_len;
1748 }
1749 Py_MEMCPY(result_s, start, end-start);
1750 return result;
1751}
1752
1753/* len(self)>=1, len(from)==len(to)==1, maxcount>=1 */
1754Py_LOCAL(PyByteArrayObject *)
1755replace_single_character_in_place(PyByteArrayObject *self,
1756 char from_c, char to_c,
1757 Py_ssize_t maxcount)
1758{
Antoine Pitroua57aae72010-06-09 16:58:35 +00001759 char *self_s, *result_s, *start, *end, *next;
1760 Py_ssize_t self_len;
1761 PyByteArrayObject *result;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001762
Antoine Pitroua57aae72010-06-09 16:58:35 +00001763 /* The result string will be the same size */
1764 self_s = PyByteArray_AS_STRING(self);
1765 self_len = PyByteArray_GET_SIZE(self);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001766
Antoine Pitroua57aae72010-06-09 16:58:35 +00001767 next = findchar(self_s, self_len, from_c);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001768
Antoine Pitroua57aae72010-06-09 16:58:35 +00001769 if (next == NULL) {
1770 /* No matches; return the original bytes */
1771 return return_self(self);
1772 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001773
Antoine Pitroua57aae72010-06-09 16:58:35 +00001774 /* Need to make a new bytes */
1775 result = (PyByteArrayObject *) PyByteArray_FromStringAndSize(NULL, self_len);
1776 if (result == NULL)
1777 return NULL;
1778 result_s = PyByteArray_AS_STRING(result);
1779 Py_MEMCPY(result_s, self_s, self_len);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001780
Antoine Pitroua57aae72010-06-09 16:58:35 +00001781 /* change everything in-place, starting with this one */
1782 start = result_s + (next-self_s);
1783 *start = to_c;
1784 start++;
1785 end = result_s + self_len;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001786
Antoine Pitroua57aae72010-06-09 16:58:35 +00001787 while (--maxcount > 0) {
1788 next = findchar(start, end-start, from_c);
1789 if (next == NULL)
1790 break;
1791 *next = to_c;
1792 start = next+1;
1793 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001794
Antoine Pitroua57aae72010-06-09 16:58:35 +00001795 return result;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00001796}
1797
1798/* len(self)>=1, len(from)==len(to)>=2, maxcount>=1 */
1799Py_LOCAL(PyByteArrayObject *)
1800replace_substring_in_place(PyByteArrayObject *self,
1801 const char *from_s, Py_ssize_t from_len,
1802 const char *to_s, Py_ssize_t to_len,
1803 Py_ssize_t maxcount)
1804{
1805 char *result_s, *start, *end;
1806 char *self_s;
1807 Py_ssize_t self_len, offset;
1808 PyByteArrayObject *result;
1809
1810 /* The result bytes will be the same size */
1811
1812 self_s = PyByteArray_AS_STRING(self);
1813 self_len = PyByteArray_GET_SIZE(self);
1814
1815 offset = findstring(self_s, self_len,
1816 from_s, from_len,
1817 0, self_len, FORWARD);
1818 if (offset == -1) {
1819 /* No matches; return the original bytes */
1820 return return_self(self);
1821 }
1822
1823 /* Need to make a new bytes */
1824 result = (PyByteArrayObject *) PyByteArray_FromStringAndSize(NULL, self_len);
1825 if (result == NULL)
1826 return NULL;
1827 result_s = PyByteArray_AS_STRING(result);
1828 Py_MEMCPY(result_s, self_s, self_len);
1829
1830 /* change everything in-place, starting with this one */
1831 start = result_s + offset;
1832 Py_MEMCPY(start, to_s, from_len);
1833 start += from_len;
1834 end = result_s + self_len;
1835
1836 while ( --maxcount > 0) {
1837 offset = findstring(start, end-start,
1838 from_s, from_len,
1839 0, end-start, FORWARD);
1840 if (offset==-1)
1841 break;
1842 Py_MEMCPY(start+offset, to_s, from_len);
1843 start += offset+from_len;
1844 }
1845
1846 return result;
1847}
1848
1849/* len(self)>=1, len(from)==1, len(to)>=2, maxcount>=1 */
1850Py_LOCAL(PyByteArrayObject *)
1851replace_single_character(PyByteArrayObject *self,
1852 char from_c,
1853 const char *to_s, Py_ssize_t to_len,
1854 Py_ssize_t maxcount)
1855{
1856 char *self_s, *result_s;
1857 char *start, *next, *end;
1858 Py_ssize_t self_len, result_len;
1859 Py_ssize_t count, product;
1860 PyByteArrayObject *result;
1861
1862 self_s = PyByteArray_AS_STRING(self);
1863 self_len = PyByteArray_GET_SIZE(self);
1864
1865 count = countchar(self_s, self_len, from_c, maxcount);
1866 if (count == 0) {
1867 /* no matches, return unchanged */
1868 return return_self(self);
1869 }
1870
1871 /* use the difference between current and new, hence the "-1" */
1872 /* result_len = self_len + count * (to_len-1) */
1873 product = count * (to_len-1);
1874 if (product / (to_len-1) != count) {
1875 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1876 return NULL;
1877 }
1878 result_len = self_len + product;
1879 if (result_len < 0) {
1880 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1881 return NULL;
1882 }
1883
1884 if ( (result = (PyByteArrayObject *)
1885 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL)
1886 return NULL;
1887 result_s = PyByteArray_AS_STRING(result);
1888
1889 start = self_s;
1890 end = self_s + self_len;
1891 while (count-- > 0) {
1892 next = findchar(start, end-start, from_c);
1893 if (next == NULL)
1894 break;
1895
1896 if (next == start) {
1897 /* replace with the 'to' */
1898 Py_MEMCPY(result_s, to_s, to_len);
1899 result_s += to_len;
1900 start += 1;
1901 } else {
1902 /* copy the unchanged old then the 'to' */
1903 Py_MEMCPY(result_s, start, next-start);
1904 result_s += (next-start);
1905 Py_MEMCPY(result_s, to_s, to_len);
1906 result_s += to_len;
1907 start = next+1;
1908 }
1909 }
1910 /* Copy the remainder of the remaining bytes */
1911 Py_MEMCPY(result_s, start, end-start);
1912
1913 return result;
1914}
1915
1916/* len(self)>=1, len(from)>=2, len(to)>=2, maxcount>=1 */
1917Py_LOCAL(PyByteArrayObject *)
1918replace_substring(PyByteArrayObject *self,
1919 const char *from_s, Py_ssize_t from_len,
1920 const char *to_s, Py_ssize_t to_len,
1921 Py_ssize_t maxcount)
1922{
1923 char *self_s, *result_s;
1924 char *start, *next, *end;
1925 Py_ssize_t self_len, result_len;
1926 Py_ssize_t count, offset, product;
1927 PyByteArrayObject *result;
1928
1929 self_s = PyByteArray_AS_STRING(self);
1930 self_len = PyByteArray_GET_SIZE(self);
1931
1932 count = countstring(self_s, self_len,
1933 from_s, from_len,
1934 0, self_len, FORWARD, maxcount);
1935 if (count == 0) {
1936 /* no matches, return unchanged */
1937 return return_self(self);
1938 }
1939
1940 /* Check for overflow */
1941 /* result_len = self_len + count * (to_len-from_len) */
1942 product = count * (to_len-from_len);
1943 if (product / (to_len-from_len) != count) {
1944 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1945 return NULL;
1946 }
1947 result_len = self_len + product;
1948 if (result_len < 0) {
1949 PyErr_SetString(PyExc_OverflowError, "replace bytes is too long");
1950 return NULL;
1951 }
1952
1953 if ( (result = (PyByteArrayObject *)
1954 PyByteArray_FromStringAndSize(NULL, result_len)) == NULL)
1955 return NULL;
1956 result_s = PyByteArray_AS_STRING(result);
1957
1958 start = self_s;
1959 end = self_s + self_len;
1960 while (count-- > 0) {
1961 offset = findstring(start, end-start,
1962 from_s, from_len,
1963 0, end-start, FORWARD);
1964 if (offset == -1)
1965 break;
1966 next = start+offset;
1967 if (next == start) {
1968 /* replace with the 'to' */
1969 Py_MEMCPY(result_s, to_s, to_len);
1970 result_s += to_len;
1971 start += from_len;
1972 } else {
1973 /* copy the unchanged old then the 'to' */
1974 Py_MEMCPY(result_s, start, next-start);
1975 result_s += (next-start);
1976 Py_MEMCPY(result_s, to_s, to_len);
1977 result_s += to_len;
1978 start = next+from_len;
1979 }
1980 }
1981 /* Copy the remainder of the remaining bytes */
1982 Py_MEMCPY(result_s, start, end-start);
1983
1984 return result;
1985}
1986
1987
1988Py_LOCAL(PyByteArrayObject *)
1989replace(PyByteArrayObject *self,
1990 const char *from_s, Py_ssize_t from_len,
1991 const char *to_s, Py_ssize_t to_len,
1992 Py_ssize_t maxcount)
1993{
1994 if (maxcount < 0) {
1995 maxcount = PY_SSIZE_T_MAX;
1996 } else if (maxcount == 0 || PyByteArray_GET_SIZE(self) == 0) {
1997 /* nothing to do; return the original bytes */
1998 return return_self(self);
1999 }
2000
2001 if (maxcount == 0 ||
2002 (from_len == 0 && to_len == 0)) {
2003 /* nothing to do; return the original bytes */
2004 return return_self(self);
2005 }
2006
2007 /* Handle zero-length special cases */
2008
2009 if (from_len == 0) {
2010 /* insert the 'to' bytes everywhere. */
2011 /* >>> "Python".replace("", ".") */
2012 /* '.P.y.t.h.o.n.' */
2013 return replace_interleave(self, to_s, to_len, maxcount);
2014 }
2015
2016 /* Except for "".replace("", "A") == "A" there is no way beyond this */
2017 /* point for an empty self bytes to generate a non-empty bytes */
2018 /* Special case so the remaining code always gets a non-empty bytes */
2019 if (PyByteArray_GET_SIZE(self) == 0) {
2020 return return_self(self);
2021 }
2022
2023 if (to_len == 0) {
Georg Brandl17cb8a82008-05-30 08:20:09 +00002024 /* delete all occurrences of 'from' bytes */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002025 if (from_len == 1) {
2026 return replace_delete_single_character(
2027 self, from_s[0], maxcount);
2028 } else {
2029 return replace_delete_substring(self, from_s, from_len, maxcount);
2030 }
2031 }
2032
2033 /* Handle special case where both bytes have the same length */
2034
2035 if (from_len == to_len) {
2036 if (from_len == 1) {
2037 return replace_single_character_in_place(
2038 self,
2039 from_s[0],
2040 to_s[0],
2041 maxcount);
2042 } else {
2043 return replace_substring_in_place(
2044 self, from_s, from_len, to_s, to_len, maxcount);
2045 }
2046 }
2047
2048 /* Otherwise use the more generic algorithms */
2049 if (from_len == 1) {
2050 return replace_single_character(self, from_s[0],
2051 to_s, to_len, maxcount);
2052 } else {
2053 /* len('from')>=2, len('to')>=1 */
2054 return replace_substring(self, from_s, from_len, to_s, to_len, maxcount);
2055 }
2056}
2057
2058
2059PyDoc_STRVAR(replace__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002060"B.replace(old, new[, count]) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002061\n\
2062Return a copy of B with all occurrences of subsection\n\
2063old replaced by new. If the optional argument count is\n\
2064given, only the first count occurrences are replaced.");
2065
2066static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002067bytearray_replace(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002068{
2069 Py_ssize_t count = -1;
2070 PyObject *from, *to, *res;
2071 Py_buffer vfrom, vto;
2072
2073 if (!PyArg_ParseTuple(args, "OO|n:replace", &from, &to, &count))
2074 return NULL;
2075
2076 if (_getbuffer(from, &vfrom) < 0)
2077 return NULL;
2078 if (_getbuffer(to, &vto) < 0) {
Martin v. Löwis423be952008-08-13 15:53:07 +00002079 PyBuffer_Release(&vfrom);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002080 return NULL;
2081 }
2082
2083 res = (PyObject *)replace((PyByteArrayObject *) self,
2084 vfrom.buf, vfrom.len,
2085 vto.buf, vto.len, count);
2086
Martin v. Löwis423be952008-08-13 15:53:07 +00002087 PyBuffer_Release(&vfrom);
2088 PyBuffer_Release(&vto);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002089 return res;
2090}
2091
2092
2093/* Overallocate the initial list to reduce the number of reallocs for small
2094 split sizes. Eg, "A A A A A A A A A A".split() (10 elements) has three
2095 resizes, to sizes 4, 8, then 16. Most observed string splits are for human
2096 text (roughly 11 words per line) and field delimited data (usually 1-10
2097 fields). For large strings the split algorithms are bandwidth limited
2098 so increasing the preallocation likely will not improve things.*/
2099
2100#define MAX_PREALLOC 12
2101
2102/* 5 splits gives 6 elements */
2103#define PREALLOC_SIZE(maxsplit) \
2104 (maxsplit >= MAX_PREALLOC ? MAX_PREALLOC : maxsplit+1)
2105
2106#define SPLIT_APPEND(data, left, right) \
2107 str = PyByteArray_FromStringAndSize((data) + (left), \
2108 (right) - (left)); \
2109 if (str == NULL) \
2110 goto onError; \
2111 if (PyList_Append(list, str)) { \
2112 Py_DECREF(str); \
2113 goto onError; \
2114 } \
2115 else \
2116 Py_DECREF(str);
2117
2118#define SPLIT_ADD(data, left, right) { \
2119 str = PyByteArray_FromStringAndSize((data) + (left), \
2120 (right) - (left)); \
2121 if (str == NULL) \
2122 goto onError; \
2123 if (count < MAX_PREALLOC) { \
2124 PyList_SET_ITEM(list, count, str); \
2125 } else { \
2126 if (PyList_Append(list, str)) { \
2127 Py_DECREF(str); \
2128 goto onError; \
2129 } \
2130 else \
2131 Py_DECREF(str); \
2132 } \
2133 count++; }
2134
2135/* Always force the list to the expected size. */
2136#define FIX_PREALLOC_SIZE(list) Py_SIZE(list) = count
2137
2138
2139Py_LOCAL_INLINE(PyObject *)
2140split_char(const char *s, Py_ssize_t len, char ch, Py_ssize_t maxcount)
2141{
2142 register Py_ssize_t i, j, count = 0;
2143 PyObject *str;
2144 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2145
2146 if (list == NULL)
2147 return NULL;
2148
2149 i = j = 0;
2150 while ((j < len) && (maxcount-- > 0)) {
2151 for(; j < len; j++) {
2152 /* I found that using memchr makes no difference */
2153 if (s[j] == ch) {
2154 SPLIT_ADD(s, i, j);
2155 i = j = j + 1;
2156 break;
2157 }
2158 }
2159 }
2160 if (i <= len) {
2161 SPLIT_ADD(s, i, len);
2162 }
2163 FIX_PREALLOC_SIZE(list);
2164 return list;
2165
2166 onError:
2167 Py_DECREF(list);
2168 return NULL;
2169}
2170
2171
2172Py_LOCAL_INLINE(PyObject *)
2173split_whitespace(const char *s, Py_ssize_t len, Py_ssize_t maxcount)
2174{
2175 register Py_ssize_t i, j, count = 0;
2176 PyObject *str;
2177 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2178
2179 if (list == NULL)
2180 return NULL;
2181
2182 for (i = j = 0; i < len; ) {
2183 /* find a token */
Eric Smith6dc46f52009-04-27 20:39:49 +00002184 while (i < len && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002185 i++;
2186 j = i;
Eric Smith6dc46f52009-04-27 20:39:49 +00002187 while (i < len && !Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002188 i++;
2189 if (j < i) {
2190 if (maxcount-- <= 0)
2191 break;
2192 SPLIT_ADD(s, j, i);
Eric Smith6dc46f52009-04-27 20:39:49 +00002193 while (i < len && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002194 i++;
2195 j = i;
2196 }
2197 }
2198 if (j < len) {
2199 SPLIT_ADD(s, j, len);
2200 }
2201 FIX_PREALLOC_SIZE(list);
2202 return list;
2203
2204 onError:
2205 Py_DECREF(list);
2206 return NULL;
2207}
2208
2209PyDoc_STRVAR(split__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002210"B.split([sep[, maxsplit]]) -> list of bytearrays\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002211\n\
2212Return a list of the sections in B, using sep as the delimiter.\n\
2213If sep is not given, B is split on ASCII whitespace characters\n\
2214(space, tab, return, newline, formfeed, vertical tab).\n\
2215If maxsplit is given, at most maxsplit splits are done.");
2216
2217static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002218bytearray_split(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002219{
2220 Py_ssize_t len = PyByteArray_GET_SIZE(self), n, i, j;
2221 Py_ssize_t maxsplit = -1, count = 0;
2222 const char *s = PyByteArray_AS_STRING(self), *sub;
2223 PyObject *list, *str, *subobj = Py_None;
2224 Py_buffer vsub;
2225#ifdef USE_FAST
2226 Py_ssize_t pos;
2227#endif
2228
2229 if (!PyArg_ParseTuple(args, "|On:split", &subobj, &maxsplit))
2230 return NULL;
2231 if (maxsplit < 0)
2232 maxsplit = PY_SSIZE_T_MAX;
2233
2234 if (subobj == Py_None)
2235 return split_whitespace(s, len, maxsplit);
2236
2237 if (_getbuffer(subobj, &vsub) < 0)
2238 return NULL;
2239 sub = vsub.buf;
2240 n = vsub.len;
2241
2242 if (n == 0) {
2243 PyErr_SetString(PyExc_ValueError, "empty separator");
Martin v. Löwis423be952008-08-13 15:53:07 +00002244 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002245 return NULL;
2246 }
Benjamin Petersonc4fe6f32008-08-19 18:57:56 +00002247 if (n == 1) {
2248 list = split_char(s, len, sub[0], maxsplit);
2249 PyBuffer_Release(&vsub);
2250 return list;
2251 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002252
2253 list = PyList_New(PREALLOC_SIZE(maxsplit));
2254 if (list == NULL) {
Martin v. Löwis423be952008-08-13 15:53:07 +00002255 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002256 return NULL;
2257 }
2258
2259#ifdef USE_FAST
2260 i = j = 0;
2261 while (maxsplit-- > 0) {
2262 pos = fastsearch(s+i, len-i, sub, n, FAST_SEARCH);
2263 if (pos < 0)
2264 break;
2265 j = i+pos;
2266 SPLIT_ADD(s, i, j);
2267 i = j + n;
2268 }
2269#else
2270 i = j = 0;
2271 while ((j+n <= len) && (maxsplit-- > 0)) {
2272 for (; j+n <= len; j++) {
2273 if (Py_STRING_MATCH(s, j, sub, n)) {
2274 SPLIT_ADD(s, i, j);
2275 i = j = j + n;
2276 break;
2277 }
2278 }
2279 }
2280#endif
2281 SPLIT_ADD(s, i, len);
2282 FIX_PREALLOC_SIZE(list);
Martin v. Löwis423be952008-08-13 15:53:07 +00002283 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002284 return list;
2285
2286 onError:
2287 Py_DECREF(list);
Martin v. Löwis423be952008-08-13 15:53:07 +00002288 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002289 return NULL;
2290}
2291
2292/* stringlib's partition shares nullbytes in some cases.
2293 undo this, we don't want the nullbytes to be shared. */
2294static PyObject *
2295make_nullbytes_unique(PyObject *result)
2296{
2297 if (result != NULL) {
2298 int i;
2299 assert(PyTuple_Check(result));
2300 assert(PyTuple_GET_SIZE(result) == 3);
2301 for (i = 0; i < 3; i++) {
2302 if (PyTuple_GET_ITEM(result, i) == (PyObject *)nullbytes) {
2303 PyObject *new = PyByteArray_FromStringAndSize(NULL, 0);
2304 if (new == NULL) {
2305 Py_DECREF(result);
2306 result = NULL;
2307 break;
2308 }
2309 Py_DECREF(nullbytes);
2310 PyTuple_SET_ITEM(result, i, new);
2311 }
2312 }
2313 }
2314 return result;
2315}
2316
2317PyDoc_STRVAR(partition__doc__,
2318"B.partition(sep) -> (head, sep, tail)\n\
2319\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002320Search for the separator sep in B, and return the part before it,\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002321the separator itself, and the part after it. If the separator is not\n\
2322found, returns B and two empty bytearray objects.");
2323
2324static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002325bytearray_partition(PyByteArrayObject *self, PyObject *sep_obj)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002326{
2327 PyObject *bytesep, *result;
2328
2329 bytesep = PyByteArray_FromObject(sep_obj);
2330 if (! bytesep)
2331 return NULL;
2332
2333 result = stringlib_partition(
2334 (PyObject*) self,
2335 PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
2336 bytesep,
2337 PyByteArray_AS_STRING(bytesep), PyByteArray_GET_SIZE(bytesep)
2338 );
2339
2340 Py_DECREF(bytesep);
2341 return make_nullbytes_unique(result);
2342}
2343
2344PyDoc_STRVAR(rpartition__doc__,
Ezio Melotti4c81fbb2010-01-25 12:02:24 +00002345"B.rpartition(sep) -> (head, sep, tail)\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002346\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002347Search for the separator sep in B, starting at the end of B,\n\
2348and return the part before it, the separator itself, and the\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002349part after it. If the separator is not found, returns two empty\n\
2350bytearray objects and B.");
2351
2352static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002353bytearray_rpartition(PyByteArrayObject *self, PyObject *sep_obj)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002354{
2355 PyObject *bytesep, *result;
2356
2357 bytesep = PyByteArray_FromObject(sep_obj);
2358 if (! bytesep)
2359 return NULL;
2360
2361 result = stringlib_rpartition(
2362 (PyObject*) self,
2363 PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
2364 bytesep,
2365 PyByteArray_AS_STRING(bytesep), PyByteArray_GET_SIZE(bytesep)
2366 );
2367
2368 Py_DECREF(bytesep);
2369 return make_nullbytes_unique(result);
2370}
2371
2372Py_LOCAL_INLINE(PyObject *)
2373rsplit_char(const char *s, Py_ssize_t len, char ch, Py_ssize_t maxcount)
2374{
2375 register Py_ssize_t i, j, count=0;
2376 PyObject *str;
2377 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2378
2379 if (list == NULL)
2380 return NULL;
2381
2382 i = j = len - 1;
2383 while ((i >= 0) && (maxcount-- > 0)) {
2384 for (; i >= 0; i--) {
2385 if (s[i] == ch) {
2386 SPLIT_ADD(s, i + 1, j + 1);
2387 j = i = i - 1;
2388 break;
2389 }
2390 }
2391 }
2392 if (j >= -1) {
2393 SPLIT_ADD(s, 0, j + 1);
2394 }
2395 FIX_PREALLOC_SIZE(list);
2396 if (PyList_Reverse(list) < 0)
2397 goto onError;
2398
2399 return list;
2400
2401 onError:
2402 Py_DECREF(list);
2403 return NULL;
2404}
2405
2406Py_LOCAL_INLINE(PyObject *)
2407rsplit_whitespace(const char *s, Py_ssize_t len, Py_ssize_t maxcount)
2408{
2409 register Py_ssize_t i, j, count = 0;
2410 PyObject *str;
2411 PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
2412
2413 if (list == NULL)
2414 return NULL;
2415
2416 for (i = j = len - 1; i >= 0; ) {
2417 /* find a token */
Eric Smith6dc46f52009-04-27 20:39:49 +00002418 while (i >= 0 && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002419 i--;
2420 j = i;
Eric Smith6dc46f52009-04-27 20:39:49 +00002421 while (i >= 0 && !Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002422 i--;
2423 if (j > i) {
2424 if (maxcount-- <= 0)
2425 break;
2426 SPLIT_ADD(s, i + 1, j + 1);
Eric Smith6dc46f52009-04-27 20:39:49 +00002427 while (i >= 0 && Py_ISSPACE(s[i]))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002428 i--;
2429 j = i;
2430 }
2431 }
2432 if (j >= 0) {
2433 SPLIT_ADD(s, 0, j + 1);
2434 }
2435 FIX_PREALLOC_SIZE(list);
2436 if (PyList_Reverse(list) < 0)
2437 goto onError;
2438
2439 return list;
2440
2441 onError:
2442 Py_DECREF(list);
2443 return NULL;
2444}
2445
2446PyDoc_STRVAR(rsplit__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002447"B.rsplit(sep[, maxsplit]) -> list of bytearrays\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002448\n\
2449Return a list of the sections in B, using sep as the delimiter,\n\
2450starting at the end of B and working to the front.\n\
2451If sep is not given, B is split on ASCII whitespace characters\n\
2452(space, tab, return, newline, formfeed, vertical tab).\n\
2453If maxsplit is given, at most maxsplit splits are done.");
2454
2455static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002456bytearray_rsplit(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002457{
2458 Py_ssize_t len = PyByteArray_GET_SIZE(self), n, i, j;
2459 Py_ssize_t maxsplit = -1, count = 0;
2460 const char *s = PyByteArray_AS_STRING(self), *sub;
2461 PyObject *list, *str, *subobj = Py_None;
2462 Py_buffer vsub;
2463
2464 if (!PyArg_ParseTuple(args, "|On:rsplit", &subobj, &maxsplit))
2465 return NULL;
2466 if (maxsplit < 0)
2467 maxsplit = PY_SSIZE_T_MAX;
2468
2469 if (subobj == Py_None)
2470 return rsplit_whitespace(s, len, maxsplit);
2471
2472 if (_getbuffer(subobj, &vsub) < 0)
2473 return NULL;
2474 sub = vsub.buf;
2475 n = vsub.len;
2476
2477 if (n == 0) {
2478 PyErr_SetString(PyExc_ValueError, "empty separator");
Martin v. Löwis423be952008-08-13 15:53:07 +00002479 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002480 return NULL;
2481 }
Benjamin Petersonc4fe6f32008-08-19 18:57:56 +00002482 else if (n == 1) {
2483 list = rsplit_char(s, len, sub[0], maxsplit);
2484 PyBuffer_Release(&vsub);
2485 return list;
2486 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002487
2488 list = PyList_New(PREALLOC_SIZE(maxsplit));
2489 if (list == NULL) {
Martin v. Löwis423be952008-08-13 15:53:07 +00002490 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002491 return NULL;
2492 }
2493
2494 j = len;
2495 i = j - n;
2496
2497 while ( (i >= 0) && (maxsplit-- > 0) ) {
2498 for (; i>=0; i--) {
2499 if (Py_STRING_MATCH(s, i, sub, n)) {
2500 SPLIT_ADD(s, i + n, j);
2501 j = i;
2502 i -= n;
2503 break;
2504 }
2505 }
2506 }
2507 SPLIT_ADD(s, 0, j);
2508 FIX_PREALLOC_SIZE(list);
2509 if (PyList_Reverse(list) < 0)
2510 goto onError;
Martin v. Löwis423be952008-08-13 15:53:07 +00002511 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002512 return list;
2513
2514onError:
2515 Py_DECREF(list);
Martin v. Löwis423be952008-08-13 15:53:07 +00002516 PyBuffer_Release(&vsub);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002517 return NULL;
2518}
2519
2520PyDoc_STRVAR(reverse__doc__,
2521"B.reverse() -> None\n\
2522\n\
2523Reverse the order of the values in B in place.");
2524static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002525bytearray_reverse(PyByteArrayObject *self, PyObject *unused)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002526{
2527 char swap, *head, *tail;
2528 Py_ssize_t i, j, n = Py_SIZE(self);
2529
2530 j = n / 2;
2531 head = self->ob_bytes;
2532 tail = head + n - 1;
2533 for (i = 0; i < j; i++) {
2534 swap = *head;
2535 *head++ = *tail;
2536 *tail-- = swap;
2537 }
2538
2539 Py_RETURN_NONE;
2540}
2541
2542PyDoc_STRVAR(insert__doc__,
2543"B.insert(index, int) -> None\n\
2544\n\
2545Insert a single item into the bytearray before the given index.");
2546static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002547bytearray_insert(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002548{
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002549 PyObject *value;
2550 int ival;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002551 Py_ssize_t where, n = Py_SIZE(self);
2552
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002553 if (!PyArg_ParseTuple(args, "nO:insert", &where, &value))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002554 return NULL;
2555
2556 if (n == PY_SSIZE_T_MAX) {
2557 PyErr_SetString(PyExc_OverflowError,
Mark Dickinson099d7a02009-09-06 10:35:38 +00002558 "cannot add more objects to bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002559 return NULL;
2560 }
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002561 if (!_getbytevalue(value, &ival))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002562 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002563 if (PyByteArray_Resize((PyObject *)self, n + 1) < 0)
2564 return NULL;
2565
2566 if (where < 0) {
2567 where += n;
2568 if (where < 0)
2569 where = 0;
2570 }
2571 if (where > n)
2572 where = n;
2573 memmove(self->ob_bytes + where + 1, self->ob_bytes + where, n - where);
Georg Brandl9a54d7c2008-07-16 23:15:30 +00002574 self->ob_bytes[where] = ival;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002575
2576 Py_RETURN_NONE;
2577}
2578
2579PyDoc_STRVAR(append__doc__,
2580"B.append(int) -> None\n\
2581\n\
2582Append a single item to the end of B.");
2583static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002584bytearray_append(PyByteArrayObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002585{
2586 int value;
2587 Py_ssize_t n = Py_SIZE(self);
2588
2589 if (! _getbytevalue(arg, &value))
2590 return NULL;
2591 if (n == PY_SSIZE_T_MAX) {
2592 PyErr_SetString(PyExc_OverflowError,
Mark Dickinson099d7a02009-09-06 10:35:38 +00002593 "cannot add more objects to bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002594 return NULL;
2595 }
2596 if (PyByteArray_Resize((PyObject *)self, n + 1) < 0)
2597 return NULL;
2598
2599 self->ob_bytes[n] = value;
2600
2601 Py_RETURN_NONE;
2602}
2603
2604PyDoc_STRVAR(extend__doc__,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002605"B.extend(iterable_of_ints) -> None\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002606\n\
2607Append all the elements from the iterator or sequence to the\n\
2608end of B.");
2609static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002610bytearray_extend(PyByteArrayObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002611{
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002612 PyObject *it, *item, *bytearray_obj;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002613 Py_ssize_t buf_size = 0, len = 0;
2614 int value;
2615 char *buf;
2616
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002617 /* bytearray_setslice code only accepts something supporting PEP 3118. */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002618 if (PyObject_CheckBuffer(arg)) {
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002619 if (bytearray_setslice(self, Py_SIZE(self), Py_SIZE(self), arg) == -1)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002620 return NULL;
2621
2622 Py_RETURN_NONE;
2623 }
2624
2625 it = PyObject_GetIter(arg);
2626 if (it == NULL)
2627 return NULL;
2628
2629 /* Try to determine the length of the argument. 32 is abitrary. */
2630 buf_size = _PyObject_LengthHint(arg, 32);
Benjamin Petersonef3e4c22009-04-11 19:48:14 +00002631 if (buf_size == -1) {
2632 Py_DECREF(it);
2633 return NULL;
2634 }
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002635
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002636 bytearray_obj = PyByteArray_FromStringAndSize(NULL, buf_size);
2637 if (bytearray_obj == NULL)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002638 return NULL;
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002639 buf = PyByteArray_AS_STRING(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002640
2641 while ((item = PyIter_Next(it)) != NULL) {
2642 if (! _getbytevalue(item, &value)) {
2643 Py_DECREF(item);
2644 Py_DECREF(it);
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002645 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002646 return NULL;
2647 }
2648 buf[len++] = value;
2649 Py_DECREF(item);
2650
2651 if (len >= buf_size) {
2652 buf_size = len + (len >> 1) + 1;
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002653 if (PyByteArray_Resize((PyObject *)bytearray_obj, buf_size) < 0) {
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002654 Py_DECREF(it);
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002655 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002656 return NULL;
2657 }
2658 /* Recompute the `buf' pointer, since the resizing operation may
2659 have invalidated it. */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002660 buf = PyByteArray_AS_STRING(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002661 }
2662 }
2663 Py_DECREF(it);
2664
2665 /* Resize down to exact size. */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002666 if (PyByteArray_Resize((PyObject *)bytearray_obj, len) < 0) {
2667 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002668 return NULL;
2669 }
2670
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002671 if (bytearray_setslice(self, Py_SIZE(self), Py_SIZE(self), bytearray_obj) == -1)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002672 return NULL;
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002673 Py_DECREF(bytearray_obj);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002674
2675 Py_RETURN_NONE;
2676}
2677
2678PyDoc_STRVAR(pop__doc__,
2679"B.pop([index]) -> int\n\
2680\n\
2681Remove and return a single item from B. If no index\n\
Benjamin Petersondcf97b92008-07-02 17:30:14 +00002682argument is given, will pop the last value.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002683static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002684bytearray_pop(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002685{
2686 int value;
2687 Py_ssize_t where = -1, n = Py_SIZE(self);
2688
2689 if (!PyArg_ParseTuple(args, "|n:pop", &where))
2690 return NULL;
2691
2692 if (n == 0) {
2693 PyErr_SetString(PyExc_OverflowError,
Mark Dickinson099d7a02009-09-06 10:35:38 +00002694 "cannot pop an empty bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002695 return NULL;
2696 }
2697 if (where < 0)
2698 where += Py_SIZE(self);
2699 if (where < 0 || where >= Py_SIZE(self)) {
2700 PyErr_SetString(PyExc_IndexError, "pop index out of range");
2701 return NULL;
2702 }
Antoine Pitrou5504e892008-12-06 21:27:53 +00002703 if (!_canresize(self))
2704 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002705
2706 value = self->ob_bytes[where];
2707 memmove(self->ob_bytes + where, self->ob_bytes + where + 1, n - where);
2708 if (PyByteArray_Resize((PyObject *)self, n - 1) < 0)
2709 return NULL;
2710
Mark Dickinson424d75a2009-09-06 10:20:23 +00002711 return PyLong_FromLong((unsigned char)value);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002712}
2713
2714PyDoc_STRVAR(remove__doc__,
2715"B.remove(int) -> None\n\
2716\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002717Remove the first occurrence of a value in B.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002718static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002719bytearray_remove(PyByteArrayObject *self, PyObject *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002720{
2721 int value;
2722 Py_ssize_t where, n = Py_SIZE(self);
2723
2724 if (! _getbytevalue(arg, &value))
2725 return NULL;
2726
2727 for (where = 0; where < n; where++) {
2728 if (self->ob_bytes[where] == value)
2729 break;
2730 }
2731 if (where == n) {
Mark Dickinson099d7a02009-09-06 10:35:38 +00002732 PyErr_SetString(PyExc_ValueError, "value not found in bytearray");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002733 return NULL;
2734 }
Antoine Pitrou5504e892008-12-06 21:27:53 +00002735 if (!_canresize(self))
2736 return NULL;
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002737
2738 memmove(self->ob_bytes + where, self->ob_bytes + where + 1, n - where);
2739 if (PyByteArray_Resize((PyObject *)self, n - 1) < 0)
2740 return NULL;
2741
2742 Py_RETURN_NONE;
2743}
2744
2745/* XXX These two helpers could be optimized if argsize == 1 */
2746
2747static Py_ssize_t
2748lstrip_helper(unsigned char *myptr, Py_ssize_t mysize,
2749 void *argptr, Py_ssize_t argsize)
2750{
2751 Py_ssize_t i = 0;
2752 while (i < mysize && memchr(argptr, myptr[i], argsize))
2753 i++;
2754 return i;
2755}
2756
2757static Py_ssize_t
2758rstrip_helper(unsigned char *myptr, Py_ssize_t mysize,
2759 void *argptr, Py_ssize_t argsize)
2760{
2761 Py_ssize_t i = mysize - 1;
2762 while (i >= 0 && memchr(argptr, myptr[i], argsize))
2763 i--;
2764 return i + 1;
2765}
2766
2767PyDoc_STRVAR(strip__doc__,
2768"B.strip([bytes]) -> bytearray\n\
2769\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002770Strip leading and trailing bytes contained in the argument\n\
2771and return the result as a new bytearray.\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002772If the argument is omitted, strip ASCII whitespace.");
2773static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002774bytearray_strip(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002775{
2776 Py_ssize_t left, right, mysize, argsize;
2777 void *myptr, *argptr;
2778 PyObject *arg = Py_None;
2779 Py_buffer varg;
2780 if (!PyArg_ParseTuple(args, "|O:strip", &arg))
2781 return NULL;
2782 if (arg == Py_None) {
2783 argptr = "\t\n\r\f\v ";
2784 argsize = 6;
2785 }
2786 else {
2787 if (_getbuffer(arg, &varg) < 0)
2788 return NULL;
2789 argptr = varg.buf;
2790 argsize = varg.len;
2791 }
2792 myptr = self->ob_bytes;
2793 mysize = Py_SIZE(self);
2794 left = lstrip_helper(myptr, mysize, argptr, argsize);
2795 if (left == mysize)
2796 right = left;
2797 else
2798 right = rstrip_helper(myptr, mysize, argptr, argsize);
2799 if (arg != Py_None)
Martin v. Löwis423be952008-08-13 15:53:07 +00002800 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002801 return PyByteArray_FromStringAndSize(self->ob_bytes + left, right - left);
2802}
2803
2804PyDoc_STRVAR(lstrip__doc__,
2805"B.lstrip([bytes]) -> bytearray\n\
2806\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002807Strip leading bytes contained in the argument\n\
2808and return the result as a new bytearray.\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002809If the argument is omitted, strip leading ASCII whitespace.");
2810static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002811bytearray_lstrip(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002812{
2813 Py_ssize_t left, right, mysize, argsize;
2814 void *myptr, *argptr;
2815 PyObject *arg = Py_None;
2816 Py_buffer varg;
2817 if (!PyArg_ParseTuple(args, "|O:lstrip", &arg))
2818 return NULL;
2819 if (arg == Py_None) {
2820 argptr = "\t\n\r\f\v ";
2821 argsize = 6;
2822 }
2823 else {
2824 if (_getbuffer(arg, &varg) < 0)
2825 return NULL;
2826 argptr = varg.buf;
2827 argsize = varg.len;
2828 }
2829 myptr = self->ob_bytes;
2830 mysize = Py_SIZE(self);
2831 left = lstrip_helper(myptr, mysize, argptr, argsize);
2832 right = mysize;
2833 if (arg != Py_None)
Martin v. Löwis423be952008-08-13 15:53:07 +00002834 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002835 return PyByteArray_FromStringAndSize(self->ob_bytes + left, right - left);
2836}
2837
2838PyDoc_STRVAR(rstrip__doc__,
2839"B.rstrip([bytes]) -> bytearray\n\
2840\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002841Strip trailing bytes contained in the argument\n\
2842and return the result as a new bytearray.\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002843If the argument is omitted, strip trailing ASCII whitespace.");
2844static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002845bytearray_rstrip(PyByteArrayObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002846{
2847 Py_ssize_t left, right, mysize, argsize;
2848 void *myptr, *argptr;
2849 PyObject *arg = Py_None;
2850 Py_buffer varg;
2851 if (!PyArg_ParseTuple(args, "|O:rstrip", &arg))
2852 return NULL;
2853 if (arg == Py_None) {
2854 argptr = "\t\n\r\f\v ";
2855 argsize = 6;
2856 }
2857 else {
2858 if (_getbuffer(arg, &varg) < 0)
2859 return NULL;
2860 argptr = varg.buf;
2861 argsize = varg.len;
2862 }
2863 myptr = self->ob_bytes;
2864 mysize = Py_SIZE(self);
2865 left = 0;
2866 right = rstrip_helper(myptr, mysize, argptr, argsize);
2867 if (arg != Py_None)
Martin v. Löwis423be952008-08-13 15:53:07 +00002868 PyBuffer_Release(&varg);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002869 return PyByteArray_FromStringAndSize(self->ob_bytes + left, right - left);
2870}
2871
2872PyDoc_STRVAR(decode_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002873"B.decode([encoding[, errors]]) -> str\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002874\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002875Decode B using the codec registered for encoding. encoding defaults\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002876to the default encoding. errors may be given to set a different error\n\
2877handling scheme. Default is 'strict' meaning that encoding errors raise\n\
2878a UnicodeDecodeError. Other possible values are 'ignore' and 'replace'\n\
2879as well as any other name registered with codecs.register_error that is\n\
2880able to handle UnicodeDecodeErrors.");
2881
2882static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002883bytearray_decode(PyObject *self, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002884{
2885 const char *encoding = NULL;
2886 const char *errors = NULL;
2887
2888 if (!PyArg_ParseTuple(args, "|ss:decode", &encoding, &errors))
2889 return NULL;
2890 if (encoding == NULL)
2891 encoding = PyUnicode_GetDefaultEncoding();
Marc-André Lemburgb2750b52008-06-06 12:18:17 +00002892 return PyUnicode_FromEncodedObject(self, encoding, errors);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002893}
2894
2895PyDoc_STRVAR(alloc_doc,
2896"B.__alloc__() -> int\n\
2897\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002898Return the number of bytes actually allocated.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002899
2900static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002901bytearray_alloc(PyByteArrayObject *self)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002902{
2903 return PyLong_FromSsize_t(self->ob_alloc);
2904}
2905
2906PyDoc_STRVAR(join_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002907"B.join(iterable_of_bytes) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002908\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00002909Concatenate any number of bytes/bytearray objects, with B\n\
2910in between each pair, and return the result as a new bytearray.");
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002911
2912static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00002913bytearray_join(PyByteArrayObject *self, PyObject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002914{
2915 PyObject *seq;
2916 Py_ssize_t mysize = Py_SIZE(self);
2917 Py_ssize_t i;
2918 Py_ssize_t n;
2919 PyObject **items;
2920 Py_ssize_t totalsize = 0;
2921 PyObject *result;
2922 char *dest;
2923
2924 seq = PySequence_Fast(it, "can only join an iterable");
2925 if (seq == NULL)
2926 return NULL;
2927 n = PySequence_Fast_GET_SIZE(seq);
2928 items = PySequence_Fast_ITEMS(seq);
2929
2930 /* Compute the total size, and check that they are all bytes */
2931 /* XXX Shouldn't we use _getbuffer() on these items instead? */
2932 for (i = 0; i < n; i++) {
2933 PyObject *obj = items[i];
2934 if (!PyByteArray_Check(obj) && !PyBytes_Check(obj)) {
2935 PyErr_Format(PyExc_TypeError,
2936 "can only join an iterable of bytes "
2937 "(item %ld has type '%.100s')",
2938 /* XXX %ld isn't right on Win64 */
2939 (long)i, Py_TYPE(obj)->tp_name);
2940 goto error;
2941 }
2942 if (i > 0)
2943 totalsize += mysize;
2944 totalsize += Py_SIZE(obj);
2945 if (totalsize < 0) {
2946 PyErr_NoMemory();
2947 goto error;
2948 }
2949 }
2950
2951 /* Allocate the result, and copy the bytes */
2952 result = PyByteArray_FromStringAndSize(NULL, totalsize);
2953 if (result == NULL)
2954 goto error;
2955 dest = PyByteArray_AS_STRING(result);
2956 for (i = 0; i < n; i++) {
2957 PyObject *obj = items[i];
2958 Py_ssize_t size = Py_SIZE(obj);
2959 char *buf;
2960 if (PyByteArray_Check(obj))
2961 buf = PyByteArray_AS_STRING(obj);
2962 else
2963 buf = PyBytes_AS_STRING(obj);
2964 if (i) {
2965 memcpy(dest, self->ob_bytes, mysize);
2966 dest += mysize;
2967 }
2968 memcpy(dest, buf, size);
2969 dest += size;
2970 }
2971
2972 /* Done */
2973 Py_DECREF(seq);
2974 return result;
2975
2976 /* Error handling */
2977 error:
2978 Py_DECREF(seq);
2979 return NULL;
2980}
2981
2982PyDoc_STRVAR(fromhex_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00002983"bytearray.fromhex(string) -> bytearray (static method)\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002984\n\
2985Create a bytearray object from a string of hexadecimal numbers.\n\
2986Spaces between two numbers are accepted.\n\
2987Example: bytearray.fromhex('B9 01EF') -> bytearray(b'\\xb9\\x01\\xef').");
2988
2989static int
2990hex_digit_to_int(Py_UNICODE c)
2991{
2992 if (c >= 128)
2993 return -1;
Eric Smith6dc46f52009-04-27 20:39:49 +00002994 if (Py_ISDIGIT(c))
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002995 return c - '0';
2996 else {
Eric Smith6dc46f52009-04-27 20:39:49 +00002997 if (Py_ISUPPER(c))
2998 c = Py_TOLOWER(c);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00002999 if (c >= 'a' && c <= 'f')
3000 return c - 'a' + 10;
3001 }
3002 return -1;
3003}
3004
3005static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003006bytearray_fromhex(PyObject *cls, PyObject *args)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003007{
3008 PyObject *newbytes, *hexobj;
3009 char *buf;
3010 Py_UNICODE *hex;
3011 Py_ssize_t hexlen, byteslen, i, j;
3012 int top, bot;
3013
3014 if (!PyArg_ParseTuple(args, "U:fromhex", &hexobj))
3015 return NULL;
3016 assert(PyUnicode_Check(hexobj));
3017 hexlen = PyUnicode_GET_SIZE(hexobj);
3018 hex = PyUnicode_AS_UNICODE(hexobj);
3019 byteslen = hexlen/2; /* This overestimates if there are spaces */
3020 newbytes = PyByteArray_FromStringAndSize(NULL, byteslen);
3021 if (!newbytes)
3022 return NULL;
3023 buf = PyByteArray_AS_STRING(newbytes);
3024 for (i = j = 0; i < hexlen; i += 2) {
3025 /* skip over spaces in the input */
3026 while (hex[i] == ' ')
3027 i++;
3028 if (i >= hexlen)
3029 break;
3030 top = hex_digit_to_int(hex[i]);
3031 bot = hex_digit_to_int(hex[i+1]);
3032 if (top == -1 || bot == -1) {
3033 PyErr_Format(PyExc_ValueError,
3034 "non-hexadecimal number found in "
3035 "fromhex() arg at position %zd", i);
3036 goto error;
3037 }
3038 buf[j++] = (top << 4) + bot;
3039 }
3040 if (PyByteArray_Resize(newbytes, j) < 0)
3041 goto error;
3042 return newbytes;
3043
3044 error:
3045 Py_DECREF(newbytes);
3046 return NULL;
3047}
3048
3049PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
3050
3051static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003052bytearray_reduce(PyByteArrayObject *self)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003053{
3054 PyObject *latin1, *dict;
3055 if (self->ob_bytes)
3056 latin1 = PyUnicode_DecodeLatin1(self->ob_bytes,
3057 Py_SIZE(self), NULL);
3058 else
3059 latin1 = PyUnicode_FromString("");
3060
3061 dict = PyObject_GetAttrString((PyObject *)self, "__dict__");
3062 if (dict == NULL) {
3063 PyErr_Clear();
3064 dict = Py_None;
3065 Py_INCREF(dict);
3066 }
3067
3068 return Py_BuildValue("(O(Ns)N)", Py_TYPE(self), latin1, "latin-1", dict);
3069}
3070
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003071PyDoc_STRVAR(sizeof_doc,
3072"B.__sizeof__() -> int\n\
3073 \n\
3074Returns the size of B in memory, in bytes");
3075static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003076bytearray_sizeof(PyByteArrayObject *self)
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003077{
Benjamin Petersonef3e4c22009-04-11 19:48:14 +00003078 Py_ssize_t res;
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003079
Benjamin Petersonef3e4c22009-04-11 19:48:14 +00003080 res = sizeof(PyByteArrayObject) + self->ob_alloc * sizeof(char);
3081 return PyLong_FromSsize_t(res);
Robert Schuppeniesfbe94c52008-07-14 10:13:31 +00003082}
3083
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003084static PySequenceMethods bytearray_as_sequence = {
3085 (lenfunc)bytearray_length, /* sq_length */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003086 (binaryfunc)PyByteArray_Concat, /* sq_concat */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003087 (ssizeargfunc)bytearray_repeat, /* sq_repeat */
3088 (ssizeargfunc)bytearray_getitem, /* sq_item */
3089 0, /* sq_slice */
3090 (ssizeobjargproc)bytearray_setitem, /* sq_ass_item */
3091 0, /* sq_ass_slice */
3092 (objobjproc)bytearray_contains, /* sq_contains */
3093 (binaryfunc)bytearray_iconcat, /* sq_inplace_concat */
3094 (ssizeargfunc)bytearray_irepeat, /* sq_inplace_repeat */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003095};
3096
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003097static PyMappingMethods bytearray_as_mapping = {
3098 (lenfunc)bytearray_length,
3099 (binaryfunc)bytearray_subscript,
3100 (objobjargproc)bytearray_ass_subscript,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003101};
3102
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003103static PyBufferProcs bytearray_as_buffer = {
3104 (getbufferproc)bytearray_getbuffer,
3105 (releasebufferproc)bytearray_releasebuffer,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003106};
3107
3108static PyMethodDef
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003109bytearray_methods[] = {
3110 {"__alloc__", (PyCFunction)bytearray_alloc, METH_NOARGS, alloc_doc},
3111 {"__reduce__", (PyCFunction)bytearray_reduce, METH_NOARGS, reduce_doc},
3112 {"__sizeof__", (PyCFunction)bytearray_sizeof, METH_NOARGS, sizeof_doc},
3113 {"append", (PyCFunction)bytearray_append, METH_O, append__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003114 {"capitalize", (PyCFunction)stringlib_capitalize, METH_NOARGS,
3115 _Py_capitalize__doc__},
3116 {"center", (PyCFunction)stringlib_center, METH_VARARGS, center__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003117 {"count", (PyCFunction)bytearray_count, METH_VARARGS, count__doc__},
3118 {"decode", (PyCFunction)bytearray_decode, METH_VARARGS, decode_doc},
3119 {"endswith", (PyCFunction)bytearray_endswith, METH_VARARGS, endswith__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003120 {"expandtabs", (PyCFunction)stringlib_expandtabs, METH_VARARGS,
3121 expandtabs__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003122 {"extend", (PyCFunction)bytearray_extend, METH_O, extend__doc__},
3123 {"find", (PyCFunction)bytearray_find, METH_VARARGS, find__doc__},
3124 {"fromhex", (PyCFunction)bytearray_fromhex, METH_VARARGS|METH_CLASS,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003125 fromhex_doc},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003126 {"index", (PyCFunction)bytearray_index, METH_VARARGS, index__doc__},
3127 {"insert", (PyCFunction)bytearray_insert, METH_VARARGS, insert__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003128 {"isalnum", (PyCFunction)stringlib_isalnum, METH_NOARGS,
3129 _Py_isalnum__doc__},
3130 {"isalpha", (PyCFunction)stringlib_isalpha, METH_NOARGS,
3131 _Py_isalpha__doc__},
3132 {"isdigit", (PyCFunction)stringlib_isdigit, METH_NOARGS,
3133 _Py_isdigit__doc__},
3134 {"islower", (PyCFunction)stringlib_islower, METH_NOARGS,
3135 _Py_islower__doc__},
3136 {"isspace", (PyCFunction)stringlib_isspace, METH_NOARGS,
3137 _Py_isspace__doc__},
3138 {"istitle", (PyCFunction)stringlib_istitle, METH_NOARGS,
3139 _Py_istitle__doc__},
3140 {"isupper", (PyCFunction)stringlib_isupper, METH_NOARGS,
3141 _Py_isupper__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003142 {"join", (PyCFunction)bytearray_join, METH_O, join_doc},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003143 {"ljust", (PyCFunction)stringlib_ljust, METH_VARARGS, ljust__doc__},
3144 {"lower", (PyCFunction)stringlib_lower, METH_NOARGS, _Py_lower__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003145 {"lstrip", (PyCFunction)bytearray_lstrip, METH_VARARGS, lstrip__doc__},
3146 {"maketrans", (PyCFunction)bytearray_maketrans, METH_VARARGS|METH_STATIC,
Georg Brandlabc38772009-04-12 15:51:51 +00003147 _Py_maketrans__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003148 {"partition", (PyCFunction)bytearray_partition, METH_O, partition__doc__},
3149 {"pop", (PyCFunction)bytearray_pop, METH_VARARGS, pop__doc__},
3150 {"remove", (PyCFunction)bytearray_remove, METH_O, remove__doc__},
3151 {"replace", (PyCFunction)bytearray_replace, METH_VARARGS, replace__doc__},
3152 {"reverse", (PyCFunction)bytearray_reverse, METH_NOARGS, reverse__doc__},
3153 {"rfind", (PyCFunction)bytearray_rfind, METH_VARARGS, rfind__doc__},
3154 {"rindex", (PyCFunction)bytearray_rindex, METH_VARARGS, rindex__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003155 {"rjust", (PyCFunction)stringlib_rjust, METH_VARARGS, rjust__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003156 {"rpartition", (PyCFunction)bytearray_rpartition, METH_O, rpartition__doc__},
3157 {"rsplit", (PyCFunction)bytearray_rsplit, METH_VARARGS, rsplit__doc__},
3158 {"rstrip", (PyCFunction)bytearray_rstrip, METH_VARARGS, rstrip__doc__},
3159 {"split", (PyCFunction)bytearray_split, METH_VARARGS, split__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003160 {"splitlines", (PyCFunction)stringlib_splitlines, METH_VARARGS,
3161 splitlines__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003162 {"startswith", (PyCFunction)bytearray_startswith, METH_VARARGS ,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003163 startswith__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003164 {"strip", (PyCFunction)bytearray_strip, METH_VARARGS, strip__doc__},
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003165 {"swapcase", (PyCFunction)stringlib_swapcase, METH_NOARGS,
3166 _Py_swapcase__doc__},
3167 {"title", (PyCFunction)stringlib_title, METH_NOARGS, _Py_title__doc__},
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003168 {"translate", (PyCFunction)bytearray_translate, METH_VARARGS,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003169 translate__doc__},
3170 {"upper", (PyCFunction)stringlib_upper, METH_NOARGS, _Py_upper__doc__},
3171 {"zfill", (PyCFunction)stringlib_zfill, METH_VARARGS, zfill__doc__},
3172 {NULL}
3173};
3174
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003175PyDoc_STRVAR(bytearray_doc,
Georg Brandl17cb8a82008-05-30 08:20:09 +00003176"bytearray(iterable_of_ints) -> bytearray\n\
3177bytearray(string, encoding[, errors]) -> bytearray\n\
3178bytearray(bytes_or_bytearray) -> mutable copy of bytes_or_bytearray\n\
3179bytearray(memory_view) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003180\n\
3181Construct an mutable bytearray object from:\n\
3182 - an iterable yielding integers in range(256)\n\
3183 - a text string encoded using the specified encoding\n\
3184 - a bytes or a bytearray object\n\
3185 - any object implementing the buffer API.\n\
3186\n\
Georg Brandl17cb8a82008-05-30 08:20:09 +00003187bytearray(int) -> bytearray\n\
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003188\n\
3189Construct a zero-initialized bytearray of the given length.");
3190
3191
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003192static PyObject *bytearray_iter(PyObject *seq);
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003193
3194PyTypeObject PyByteArray_Type = {
3195 PyVarObject_HEAD_INIT(&PyType_Type, 0)
3196 "bytearray",
3197 sizeof(PyByteArrayObject),
3198 0,
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003199 (destructor)bytearray_dealloc, /* tp_dealloc */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003200 0, /* tp_print */
3201 0, /* tp_getattr */
3202 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003203 0, /* tp_reserved */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003204 (reprfunc)bytearray_repr, /* tp_repr */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003205 0, /* tp_as_number */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003206 &bytearray_as_sequence, /* tp_as_sequence */
3207 &bytearray_as_mapping, /* tp_as_mapping */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003208 0, /* tp_hash */
3209 0, /* tp_call */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003210 bytearray_str, /* tp_str */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003211 PyObject_GenericGetAttr, /* tp_getattro */
3212 0, /* tp_setattro */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003213 &bytearray_as_buffer, /* tp_as_buffer */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003214 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003215 bytearray_doc, /* tp_doc */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003216 0, /* tp_traverse */
3217 0, /* tp_clear */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003218 (richcmpfunc)bytearray_richcompare, /* tp_richcompare */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003219 0, /* tp_weaklistoffset */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003220 bytearray_iter, /* tp_iter */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003221 0, /* tp_iternext */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003222 bytearray_methods, /* tp_methods */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003223 0, /* tp_members */
3224 0, /* tp_getset */
3225 0, /* tp_base */
3226 0, /* tp_dict */
3227 0, /* tp_descr_get */
3228 0, /* tp_descr_set */
3229 0, /* tp_dictoffset */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003230 (initproc)bytearray_init, /* tp_init */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003231 PyType_GenericAlloc, /* tp_alloc */
3232 PyType_GenericNew, /* tp_new */
3233 PyObject_Del, /* tp_free */
3234};
3235
3236/*********************** Bytes Iterator ****************************/
3237
3238typedef struct {
3239 PyObject_HEAD
3240 Py_ssize_t it_index;
3241 PyByteArrayObject *it_seq; /* Set to NULL when iterator is exhausted */
3242} bytesiterobject;
3243
3244static void
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003245bytearrayiter_dealloc(bytesiterobject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003246{
3247 _PyObject_GC_UNTRACK(it);
3248 Py_XDECREF(it->it_seq);
3249 PyObject_GC_Del(it);
3250}
3251
3252static int
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003253bytearrayiter_traverse(bytesiterobject *it, visitproc visit, void *arg)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003254{
3255 Py_VISIT(it->it_seq);
3256 return 0;
3257}
3258
3259static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003260bytearrayiter_next(bytesiterobject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003261{
3262 PyByteArrayObject *seq;
3263 PyObject *item;
3264
3265 assert(it != NULL);
3266 seq = it->it_seq;
3267 if (seq == NULL)
3268 return NULL;
3269 assert(PyByteArray_Check(seq));
3270
3271 if (it->it_index < PyByteArray_GET_SIZE(seq)) {
3272 item = PyLong_FromLong(
3273 (unsigned char)seq->ob_bytes[it->it_index]);
3274 if (item != NULL)
3275 ++it->it_index;
3276 return item;
3277 }
3278
3279 Py_DECREF(seq);
3280 it->it_seq = NULL;
3281 return NULL;
3282}
3283
3284static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003285bytesarrayiter_length_hint(bytesiterobject *it)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003286{
3287 Py_ssize_t len = 0;
3288 if (it->it_seq)
3289 len = PyByteArray_GET_SIZE(it->it_seq) - it->it_index;
3290 return PyLong_FromSsize_t(len);
3291}
3292
3293PyDoc_STRVAR(length_hint_doc,
3294 "Private method returning an estimate of len(list(it)).");
3295
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003296static PyMethodDef bytearrayiter_methods[] = {
3297 {"__length_hint__", (PyCFunction)bytesarrayiter_length_hint, METH_NOARGS,
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003298 length_hint_doc},
3299 {NULL, NULL} /* sentinel */
3300};
3301
3302PyTypeObject PyByteArrayIter_Type = {
3303 PyVarObject_HEAD_INIT(&PyType_Type, 0)
3304 "bytearray_iterator", /* tp_name */
3305 sizeof(bytesiterobject), /* tp_basicsize */
3306 0, /* tp_itemsize */
3307 /* methods */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003308 (destructor)bytearrayiter_dealloc, /* tp_dealloc */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003309 0, /* tp_print */
3310 0, /* tp_getattr */
3311 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003312 0, /* tp_reserved */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003313 0, /* tp_repr */
3314 0, /* tp_as_number */
3315 0, /* tp_as_sequence */
3316 0, /* tp_as_mapping */
3317 0, /* tp_hash */
3318 0, /* tp_call */
3319 0, /* tp_str */
3320 PyObject_GenericGetAttr, /* tp_getattro */
3321 0, /* tp_setattro */
3322 0, /* tp_as_buffer */
3323 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
3324 0, /* tp_doc */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003325 (traverseproc)bytearrayiter_traverse, /* tp_traverse */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003326 0, /* tp_clear */
3327 0, /* tp_richcompare */
3328 0, /* tp_weaklistoffset */
3329 PyObject_SelfIter, /* tp_iter */
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003330 (iternextfunc)bytearrayiter_next, /* tp_iternext */
3331 bytearrayiter_methods, /* tp_methods */
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003332 0,
3333};
3334
3335static PyObject *
Benjamin Peterson153c70f2009-04-18 15:42:12 +00003336bytearray_iter(PyObject *seq)
Christian Heimes2c9c7a52008-05-26 13:42:13 +00003337{
3338 bytesiterobject *it;
3339
3340 if (!PyByteArray_Check(seq)) {
3341 PyErr_BadInternalCall();
3342 return NULL;
3343 }
3344 it = PyObject_GC_New(bytesiterobject, &PyByteArrayIter_Type);
3345 if (it == NULL)
3346 return NULL;
3347 it->it_index = 0;
3348 Py_INCREF(seq);
3349 it->it_seq = (PyByteArrayObject *)seq;
3350 _PyObject_GC_TRACK(it);
3351 return (PyObject *)it;
3352}