diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py index 9794cd058d480..f71fa9e010cfc 100644 --- a/mypyc/irbuild/expression.py +++ b/mypyc/irbuild/expression.py @@ -27,6 +27,7 @@ from mypyc.irbuild.format_str_tokenizer import ( tokenizer_printf_style, join_formatted_strings, convert_expr ) +from mypyc.primitives.bytes_ops import bytes_slice_op from mypyc.primitives.registry import CFunctionDescription, builtin_names, binary_ops from mypyc.primitives.generic_ops import iter_op from mypyc.primitives.misc_ops import new_slice_op, ellipsis_op, type_op, get_module_dict_op @@ -441,7 +442,7 @@ def try_gen_slice_op(builder: IRBuilder, base: Value, index: SliceExpr) -> Optio # Replace missing end index with the largest short integer # (a sequence can't be longer). end = builder.load_int(MAX_SHORT_INT) - candidates = [list_slice_op, tuple_slice_op, str_slice_op] + candidates = [list_slice_op, tuple_slice_op, str_slice_op, bytes_slice_op] return builder.builder.matching_call_c(candidates, [base, begin, end], index.line) return None diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h index b5681e346a00c..9ef2085f598d5 100644 --- a/mypyc/lib-rt/CPy.h +++ b/mypyc/lib-rt/CPy.h @@ -402,6 +402,7 @@ PyObject *CPy_Encode(PyObject *obj, PyObject *encoding, PyObject *errors); // Bytes operations +PyObject *CPyBytes_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end); 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 e8efe12e011b1..5d855805945bf 100644 --- a/mypyc/lib-rt/bytes_ops.c +++ b/mypyc/lib-rt/bytes_ops.c @@ -42,6 +42,34 @@ PyObject *CPyBytes_Concat(PyObject *a, PyObject *b) { } } +static inline Py_ssize_t Clamp(Py_ssize_t a, Py_ssize_t b, Py_ssize_t c) { + return a < b ? b : (a >= c ? c : a); +} + +PyObject *CPyBytes_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) { + if ((PyBytes_Check(obj) || PyByteArray_Check(obj)) + && CPyTagged_CheckShort(start) && CPyTagged_CheckShort(end)) { + Py_ssize_t startn = CPyTagged_ShortAsSsize_t(start); + Py_ssize_t endn = CPyTagged_ShortAsSsize_t(end); + Py_ssize_t len = ((PyVarObject *)obj)->ob_size; + if (startn < 0) { + startn += len; + } + if (endn < 0) { + endn += len; + } + startn = Clamp(startn, 0, len); + endn = Clamp(endn, 0, len); + Py_ssize_t slice_len = endn - startn; + if (PyBytes_Check(obj)) { + return PyBytes_FromStringAndSize(PyBytes_AS_STRING(obj) + startn, slice_len); + } else { + return PyByteArray_FromStringAndSize(PyByteArray_AS_STRING(obj) + startn, slice_len); + } + } + return CPyObject_GetSlice(obj, start, end); +} + // Like _PyBytes_Join but fallback to dynamic call if 'sep' is not bytes // (mostly commonly, for bytearrays) PyObject *CPyBytes_Join(PyObject *sep, PyObject *iter) { diff --git a/mypyc/primitives/bytes_ops.py b/mypyc/primitives/bytes_ops.py index 3fa0300cad255..a276951a0c2de 100644 --- a/mypyc/primitives/bytes_ops.py +++ b/mypyc/primitives/bytes_ops.py @@ -6,7 +6,7 @@ str_rprimitive, RUnion, int_rprimitive ) from mypyc.primitives.registry import ( - load_address_op, function_op, method_op, binary_op + load_address_op, function_op, method_op, binary_op, custom_op ) # Get the 'bytes' type object. @@ -41,6 +41,13 @@ error_kind=ERR_MAGIC, steals=[True, False]) +# bytes[begin:end] +bytes_slice_op = custom_op( + arg_types=[bytes_rprimitive, int_rprimitive, int_rprimitive], + return_type=bytes_rprimitive, + c_function_name='CPyBytes_GetSlice', + error_kind=ERR_MAGIC) + # bytes[index] # bytearray[index] method_op( diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py index c5c1945c2c985..362495fdaa683 100644 --- a/mypyc/test-data/fixtures/ir.py +++ b/mypyc/test-data/fixtures/ir.py @@ -104,7 +104,10 @@ def __init__(self, x: object) -> None: ... def __add__(self, x: bytes) -> bytes: ... def __eq__(self, x: object) -> bool: ... def __ne__(self, x: object) -> bool: ... + @overload def __getitem__(self, i: int) -> int: ... + @overload + def __getitem__(self, i: slice) -> bytes: ... def join(self, x: Iterable[object]) -> bytes: ... def decode(self, x: str=..., y: str=...) -> str: ... diff --git a/mypyc/test-data/irbuild-bytes.test b/mypyc/test-data/irbuild-bytes.test index 630e1b2441d12..19f8893199dcc 100644 --- a/mypyc/test-data/irbuild-bytes.test +++ b/mypyc/test-data/irbuild-bytes.test @@ -62,6 +62,18 @@ L0: c = r6 return 1 +[case testBytesSlicing] +def f(a: bytes, start: int, end: int) -> bytes: + return a[start:end] +[out] +def f(a, start, end): + a :: bytes + start, end :: int + r0 :: bytes +L0: + r0 = CPyBytes_GetSlice(a, start, end) + return r0 + [case testBytesIndex] def f(a: bytes, i: int) -> int: return a[i] diff --git a/mypyc/test-data/run-bytes.test b/mypyc/test-data/run-bytes.test index 46402fe8fe253..be8b62a8ec376 100644 --- a/mypyc/test-data/run-bytes.test +++ b/mypyc/test-data/run-bytes.test @@ -109,9 +109,41 @@ def test_len() -> None: assert len(b) == 3 assert len(bytes()) == 0 +[case testBytesSlicing] +def test_bytes_slicing() -> None: + b = b'abcdefg' + zero = int() + ten = 10 + zero + two = 2 + zero + five = 5 + zero + seven = 7 + zero + assert b[:ten] == b'abcdefg' + assert b[0:seven] == b'abcdefg' + assert b[0:(len(b)+1)] == b'abcdefg' + assert b[two:five] == b'cde' + assert b[two:two] == b'' + assert b[-two:-two] == b'' + assert b[-ten:(-ten+1)] == b'' + assert b[:-two] == b'abcde' + assert b[:two] == b'ab' + assert b[:] == b'abcdefg' + assert b[-two:] == b'fg' + assert b[zero:] == b'abcdefg' + assert b[:zero] == b'' + assert b[-ten:] == b'abcdefg' + assert b[-ten:ten] == b'abcdefg' + big_ints = [1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000, 2**24, 2**63] + for big_int in big_ints: + assert b[1:big_int] == b'bcdefg' + assert b[big_int:] == b'' + assert b[-big_int:-1] == b'abcdef' + assert b[-big_int:big_int] == b'abcdefg' + assert type(b[-big_int:-1]) == bytes + assert type(b[-ten:]) == bytes + assert type(b[:]) == bytes + [case testBytearrayBasics] from typing import Any -from testutil import assertRaises def test_basics() -> None: brr1: bytes = bytearray(3) @@ -138,6 +170,41 @@ def test_bytearray_passed_into_bytes() -> None: brr1: Any = bytearray() assert f(brr1) +[case testBytearraySlicing] +def test_bytearray_slicing() -> None: + b: bytes = bytearray(b'abcdefg') + zero = int() + ten = 10 + zero + two = 2 + zero + five = 5 + zero + seven = 7 + zero + assert b[:ten] == b'abcdefg' + assert b[0:seven] == b'abcdefg' + assert b[two:five] == b'cde' + assert b[two:two] == b'' + assert b[-two:-two] == b'' + assert b[-ten:(-ten+1)] == b'' + assert b[:-two] == b'abcde' + assert b[:two] == b'ab' + assert b[:] == b'abcdefg' + assert b[-two:] == b'fg' + assert b[zero:] == b'abcdefg' + assert b[:zero] == b'' + assert b[-ten:] == b'abcdefg' + assert b[-ten:ten] == b'abcdefg' + big_ints = [1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000, 2**24, 2**63] + for big_int in big_ints: + assert b[1:big_int] == b'bcdefg' + assert b[big_int:] == b'' + assert b[-big_int:-1] == b'abcdef' + assert b[-big_int:big_int] == b'abcdefg' + assert type(b[-big_int:-1]) == bytearray + assert type(b[-ten:]) == bytearray + assert type(b[:]) == bytearray + +[case testBytearrayIndexing] +from testutil import assertRaises + def test_bytearray_indexing() -> None: b: bytes = bytearray(b'\xae\x80\xfe\x15') assert b[0] == 174