python · JukkaL · Aug 12, 2021 · Aug 10, 2021 · Aug 10, 2021 · Aug 11, 2021
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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