diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h index 739e46057d5fc..fdbf5e1e59edc 100644 --- a/mypyc/lib-rt/CPy.h +++ b/mypyc/lib-rt/CPy.h @@ -400,6 +400,7 @@ Py_ssize_t CPyStr_Size_size_t(PyObject *str); // Bytes operations +CPyTagged CPyBytes_GetItem(PyObject *o, CPyTagged index); PyObject *CPyBytes_Concat(PyObject *a, PyObject *b); PyObject *CPyBytes_Join(PyObject *sep, PyObject *iter); diff --git a/mypyc/lib-rt/bytes_ops.c b/mypyc/lib-rt/bytes_ops.c index 9147dcbe0ec11..abfd68c3bd850 100644 --- a/mypyc/lib-rt/bytes_ops.c +++ b/mypyc/lib-rt/bytes_ops.c @@ -5,6 +5,25 @@ #include #include "CPy.h" +CPyTagged CPyBytes_GetItem(PyObject *o, CPyTagged index) { + if (CPyTagged_CheckShort(index)) { + Py_ssize_t n = CPyTagged_ShortAsSsize_t(index); + Py_ssize_t size = ((PyVarObject *)o)->ob_size; + if (n < 0) + n += size; + if (n < 0 || n >= size) { + PyErr_SetString(PyExc_IndexError, "index out of range"); + return CPY_INT_TAG; + } + unsigned char num = PyBytes_Check(o) ? ((PyBytesObject *)o)->ob_sval[n] + : ((PyByteArrayObject *)o)->ob_bytes[n]; + return num << 1; + } else { + PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG); + return CPY_INT_TAG; + } +} + PyObject *CPyBytes_Concat(PyObject *a, PyObject *b) { if (PyBytes_Check(a) && PyBytes_Check(b)) { Py_ssize_t a_len = ((PyVarObject *)a)->ob_size; diff --git a/mypyc/primitives/bytes_ops.py b/mypyc/primitives/bytes_ops.py index a95a7da3fa190..3fa0300cad255 100644 --- a/mypyc/primitives/bytes_ops.py +++ b/mypyc/primitives/bytes_ops.py @@ -3,7 +3,7 @@ from mypyc.ir.ops import ERR_MAGIC from mypyc.ir.rtypes import ( object_rprimitive, bytes_rprimitive, list_rprimitive, dict_rprimitive, - str_rprimitive, RUnion + str_rprimitive, RUnion, int_rprimitive ) from mypyc.primitives.registry import ( load_address_op, function_op, method_op, binary_op @@ -41,11 +41,19 @@ error_kind=ERR_MAGIC, steals=[True, False]) +# bytes[index] +# bytearray[index] +method_op( + name='__getitem__', + arg_types=[bytes_rprimitive, int_rprimitive], + return_type=int_rprimitive, + c_function_name='CPyBytes_GetItem', + error_kind=ERR_MAGIC) + # bytes.join(obj) method_op( name='join', arg_types=[bytes_rprimitive, object_rprimitive], return_type=bytes_rprimitive, c_function_name='CPyBytes_Join', - error_kind=ERR_MAGIC -) + error_kind=ERR_MAGIC) diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py index f64db578c868b..7fe7d42173977 100644 --- a/mypyc/test-data/fixtures/ir.py +++ b/mypyc/test-data/fixtures/ir.py @@ -116,6 +116,8 @@ def __init__(self, x: object) -> None: pass @overload def __init__(self, string: str, encoding: str, err: str = ...) -> None: pass def __add__(self, s: bytes) -> bytearray: ... + def __setitem__(self, i: int, o: int) -> None: ... + def __getitem__(self, i: int) -> int: ... class bool(int): def __init__(self, o: object = ...) -> None: ... diff --git a/mypyc/test-data/irbuild-bytes.test b/mypyc/test-data/irbuild-bytes.test index d86dc3227b4a3..630e1b2441d12 100644 --- a/mypyc/test-data/irbuild-bytes.test +++ b/mypyc/test-data/irbuild-bytes.test @@ -62,11 +62,22 @@ L0: c = r6 return 1 +[case testBytesIndex] +def f(a: bytes, i: int) -> int: + return a[i] +[out] +def f(a, i): + a :: bytes + i, r0 :: int +L0: + r0 = CPyBytes_GetItem(a, i) + return r0 + [case testBytesConcat] -def f1(a: bytes, b: bytes) -> bytes: +def f(a: bytes, b: bytes) -> bytes: return a + b [out] -def f1(a, b): +def f(a, b): a, b, r0 :: bytes L0: r0 = CPyBytes_Concat(a, b) diff --git a/mypyc/test-data/run-bytes.test b/mypyc/test-data/run-bytes.test index 01fd6c45c3f9d..46402fe8fe253 100644 --- a/mypyc/test-data/run-bytes.test +++ b/mypyc/test-data/run-bytes.test @@ -43,6 +43,8 @@ def test_bytes_init() -> None: pass [case testBytesOps] +from testutil import assertRaises + def test_indexing() -> None: # Use bytes() to avoid constant folding b = b'asdf' + bytes() @@ -50,9 +52,18 @@ def test_indexing() -> None: assert b[1] == 115 assert b[3] == 102 assert b[-1] == 102 - b = b'\xfe\x15' + bytes() - assert b[0] == 254 - assert b[1] == 21 + b = b'\xae\x80\xfe\x15' + bytes() + assert b[0] == 174 + assert b[1] == 128 + assert b[2] == 254 + assert b[3] == 21 + assert b[-4] == 174 + with assertRaises(IndexError, "index out of range"): + b[4] + with assertRaises(IndexError, "index out of range"): + b[-5] + with assertRaises(IndexError, "index out of range"): + b[2**26] def test_concat() -> None: b1 = b'123' + bytes() @@ -100,6 +111,7 @@ def test_len() -> None: [case testBytearrayBasics] from typing import Any +from testutil import assertRaises def test_basics() -> None: brr1: bytes = bytearray(3) @@ -126,6 +138,27 @@ def test_bytearray_passed_into_bytes() -> None: brr1: Any = bytearray() assert f(brr1) +def test_bytearray_indexing() -> None: + b: bytes = bytearray(b'\xae\x80\xfe\x15') + assert b[0] == 174 + assert b[1] == 128 + assert b[2] == 254 + assert b[3] == 21 + assert b[-4] == 174 + with assertRaises(IndexError, "index out of range"): + b[4] + with assertRaises(IndexError, "index out of range"): + b[-5] + b2 = bytearray([175, 255, 128, 22]) + assert b2[0] == 175 + assert b2[1] == 255 + assert b2[-1] == 22 + assert b2[2] == 128 + with assertRaises(ValueError, "byte must be in range(0, 256)"): + b2[0] = -1 + with assertRaises(ValueError, "byte must be in range(0, 256)"): + b2[0] = 256 + [case testBytesJoin] from typing import Any from testutil import assertRaises