Issue #12170: The count(), find(), rfind(), index() and rindex() methods
of bytes and bytearray objects now accept an integer between 0 and 255
as their first argument. Patch by Petri Lehtinen.
diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
index 0c9add0..6f0b28d 100644
--- a/Objects/bytearrayobject.c
+++ b/Objects/bytearrayobject.c
@@ -1071,24 +1071,41 @@
bytearray_find_internal(PyByteArrayObject *self, PyObject *args, int dir)
{
PyObject *subobj;
+ char byte;
Py_buffer subbuf;
+ const char *sub;
+ Py_ssize_t sub_len;
Py_ssize_t start=0, end=PY_SSIZE_T_MAX;
Py_ssize_t res;
- if (!stringlib_parse_args_finds("find/rfind/index/rindex",
- args, &subobj, &start, &end))
+ if (!stringlib_parse_args_finds_byte("find/rfind/index/rindex",
+ args, &subobj, &byte, &start, &end))
return -2;
- if (_getbuffer(subobj, &subbuf) < 0)
- return -2;
+
+ if (subobj) {
+ if (_getbuffer(subobj, &subbuf) < 0)
+ return -2;
+
+ sub = subbuf.buf;
+ sub_len = subbuf.len;
+ }
+ else {
+ sub = &byte;
+ sub_len = 1;
+ }
+
if (dir > 0)
res = stringlib_find_slice(
PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
- subbuf.buf, subbuf.len, start, end);
+ sub, sub_len, start, end);
else
res = stringlib_rfind_slice(
PyByteArray_AS_STRING(self), PyByteArray_GET_SIZE(self),
- subbuf.buf, subbuf.len, start, end);
- PyBuffer_Release(&subbuf);
+ sub, sub_len, start, end);
+
+ if (subobj)
+ PyBuffer_Release(&subbuf);
+
return res;
}
@@ -1121,23 +1138,39 @@
bytearray_count(PyByteArrayObject *self, PyObject *args)
{
PyObject *sub_obj;
- const char *str = PyByteArray_AS_STRING(self);
+ const char *str = PyByteArray_AS_STRING(self), *sub;
+ Py_ssize_t sub_len;
+ char byte;
Py_ssize_t start = 0, end = PY_SSIZE_T_MAX;
+
Py_buffer vsub;
PyObject *count_obj;
- if (!stringlib_parse_args_finds("count", args, &sub_obj, &start, &end))
+ if (!stringlib_parse_args_finds_byte("count", args, &sub_obj, &byte,
+ &start, &end))
return NULL;
- if (_getbuffer(sub_obj, &vsub) < 0)
- return NULL;
+ if (sub_obj) {
+ if (_getbuffer(sub_obj, &vsub) < 0)
+ return NULL;
+
+ sub = vsub.buf;
+ sub_len = vsub.len;
+ }
+ else {
+ sub = &byte;
+ sub_len = 1;
+ }
ADJUST_INDICES(start, end, PyByteArray_GET_SIZE(self));
count_obj = PyLong_FromSsize_t(
- stringlib_count(str + start, end - start, vsub.buf, vsub.len, PY_SSIZE_T_MAX)
+ stringlib_count(str + start, end - start, sub, sub_len, PY_SSIZE_T_MAX)
);
- PyBuffer_Release(&vsub);
+
+ if (sub_obj)
+ PyBuffer_Release(&vsub);
+
return count_obj;
}
diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c
index 17e31b9..7438a70 100644
--- a/Objects/bytesobject.c
+++ b/Objects/bytesobject.c
@@ -1230,31 +1230,42 @@
bytes_find_internal(PyBytesObject *self, PyObject *args, int dir)
{
PyObject *subobj;
+ char byte;
+ Py_buffer subbuf;
const char *sub;
Py_ssize_t sub_len;
Py_ssize_t start=0, end=PY_SSIZE_T_MAX;
+ Py_ssize_t res;
- if (!stringlib_parse_args_finds("find/rfind/index/rindex",
- args, &subobj, &start, &end))
+ if (!stringlib_parse_args_finds_byte("find/rfind/index/rindex",
+ args, &subobj, &byte, &start, &end))
return -2;
- if (PyBytes_Check(subobj)) {
- sub = PyBytes_AS_STRING(subobj);
- sub_len = PyBytes_GET_SIZE(subobj);
+ if (subobj) {
+ if (_getbuffer(subobj, &subbuf) < 0)
+ return -2;
+
+ sub = subbuf.buf;
+ sub_len = subbuf.len;
}
- else if (PyObject_AsCharBuffer(subobj, &sub, &sub_len))
- /* XXX - the "expected a character buffer object" is pretty
- confusing for a non-expert. remap to something else ? */
- return -2;
+ else {
+ sub = &byte;
+ sub_len = 1;
+ }
if (dir > 0)
- return stringlib_find_slice(
+ res = stringlib_find_slice(
PyBytes_AS_STRING(self), PyBytes_GET_SIZE(self),
sub, sub_len, start, end);
else
- return stringlib_rfind_slice(
+ res = stringlib_rfind_slice(
PyBytes_AS_STRING(self), PyBytes_GET_SIZE(self),
sub, sub_len, start, end);
+
+ if (subobj)
+ PyBuffer_Release(&subbuf);
+
+ return res;
}
@@ -1480,23 +1491,38 @@
PyObject *sub_obj;
const char *str = PyBytes_AS_STRING(self), *sub;
Py_ssize_t sub_len;
+ char byte;
Py_ssize_t start = 0, end = PY_SSIZE_T_MAX;
- if (!stringlib_parse_args_finds("count", args, &sub_obj, &start, &end))
+ Py_buffer vsub;
+ PyObject *count_obj;
+
+ if (!stringlib_parse_args_finds_byte("count", args, &sub_obj, &byte,
+ &start, &end))
return NULL;
- if (PyBytes_Check(sub_obj)) {
- sub = PyBytes_AS_STRING(sub_obj);
- sub_len = PyBytes_GET_SIZE(sub_obj);
+ if (sub_obj) {
+ if (_getbuffer(sub_obj, &vsub) < 0)
+ return NULL;
+
+ sub = vsub.buf;
+ sub_len = vsub.len;
}
- else if (PyObject_AsCharBuffer(sub_obj, &sub, &sub_len))
- return NULL;
+ else {
+ sub = &byte;
+ sub_len = 1;
+ }
ADJUST_INDICES(start, end, PyBytes_GET_SIZE(self));
- return PyLong_FromSsize_t(
+ count_obj = PyLong_FromSsize_t(
stringlib_count(str + start, end - start, sub, sub_len, PY_SSIZE_T_MAX)
);
+
+ if (sub_obj)
+ PyBuffer_Release(&vsub);
+
+ return count_obj;
}
diff --git a/Objects/stringlib/find.h b/Objects/stringlib/find.h
index 7cce156..00eaf1b 100644
--- a/Objects/stringlib/find.h
+++ b/Objects/stringlib/find.h
@@ -167,4 +167,47 @@
return 0;
}
+#else /* !STRINGLIB_IS_UNICODE */
+
+/*
+Wraps stringlib_parse_args_finds() and additionally checks whether the
+first argument is an integer in range(0, 256).
+
+If this is the case, writes the integer value to the byte parameter
+and sets subobj to NULL. Otherwise, sets the first argument to subobj
+and doesn't touch byte. The other parameters are similar to those of
+stringlib_parse_args_finds().
+*/
+
+Py_LOCAL_INLINE(int)
+STRINGLIB(parse_args_finds_byte)(const char *function_name, PyObject *args,
+ PyObject **subobj, char *byte,
+ Py_ssize_t *start, Py_ssize_t *end)
+{
+ PyObject *tmp_subobj;
+ Py_ssize_t ival;
+
+ if(!STRINGLIB(parse_args_finds)(function_name, args, &tmp_subobj,
+ start, end))
+ return 0;
+
+ ival = PyNumber_AsSsize_t(tmp_subobj, PyExc_ValueError);
+ if (ival == -1 && PyErr_Occurred()) {
+ PyErr_Clear();
+ *subobj = tmp_subobj;
+ }
+ else {
+ /* The first argument was an integer */
+ if(ival < 0 || ival > 255) {
+ PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
+ return 0;
+ }
+
+ *subobj = NULL;
+ *byte = (char)ival;
+ }
+
+ return 1;
+}
+
#endif /* STRINGLIB_IS_UNICODE */