python · JukkaL · Aug 11, 2021 · Aug 9, 2021 · Aug 9, 2021 · Aug 9, 2021
diff --git 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
@@ -5,6 +5,25 @@
 #include <Python.h>
 #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
@@ -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
@@ -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
@@ -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
@@ -43,16 +43,27 @@ 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()
     assert b[0] == 97
     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