Add bytearray wrapper type (#654)

* Add bytearray wrapper type

Author: noahbkim

docs/api_core.rst

@@ -1028,7 +1028,7 @@ Wrapper classes
 
    .. cpp:function:: bytes(handle h)
 
-      Performs a cast within Python. This is equivalent equivalent to
+      Performs a cast within Python. This is equivalent to
       the Python expression ``bytes(h)``.
 
    .. cpp:function:: bytes(const char * s)
@@ -1052,6 +1052,42 @@ Wrapper classes
       Convert a Python ``bytes`` object into a byte buffer of length :cpp:func:`bytes::size()` bytes.
 
 
+.. cpp:class:: bytearray: public object
+
+   This wrapper class represents Python ``bytearray`` instances.
+
+   .. cpp:function:: bytearray()
+
+      Create an empty ``bytearray``.
+
+   .. cpp:function:: bytearray(handle h)
+
+      Performs a cast within Python. This is equivalent to
+      the Python expression ``bytearray(h)``.
+
+   .. cpp:function:: bytearray(const void * buf, size_t n)
+
+      Convert a byte buffer ``buf`` of length ``n`` bytes into a Python ``bytearray`` object.  The buffer can contain embedded null bytes.
+
+   .. cpp:function:: const char * c_str() const
+
+      Convert a Python ``bytearray`` object into a null-terminated C-style string.
+
+   .. cpp:function:: size_t size() const
+
+      Return the size in bytes.
+
+   .. cpp:function:: const void * data() const
+
+      Convert a Python ``bytearray`` object into a byte buffer of length :cpp:func:`bytearray::size()` bytes.
+
+   .. cpp:function:: void resize(size_t n)
+
+      Resize the internal buffer of a Python ``bytearray`` object to ``n``. Any
+      space added by this method, which calls `PyByteArray_Resize`, will not be
+      initialized and may contain random data.
+
+
 .. cpp:class:: type_object: public object
 
    Wrapper class representing Python ``type`` instances.

docs/changelog.rst

@@ -15,10 +15,17 @@ case, both modules must use the same nanobind ABI version, or they will be
 isolated from each other. Releases that don't explicitly mention an ABI version
 below inherit that of the preceding release.
 
+Version 2.1.1 (TBA)
+-------------------
+
+* Added the :cpp:class:`bytearray` wrapper type. (PR `#654
+  <https://github.com/wjakob/nanobind/pull/654>`__)
+
+
 Version 2.1.0 (Aug 11, 2024)
 ----------------------------
 
-* Temporary workaround for a internal compiler error in version 17.10 the MSVC
+* Temporary workaround for a internal compiler error in version 17.10 of the MSVC
   compiler. This workaround will be removed once fixed versions are deployed on
   GitHub actions. (issue `#613
   <https://github.com/wjakob/nanobind/issues/613>`__, commit `f2438b

docs/exchanging.rst

@@ -389,7 +389,7 @@ multithreaded computations.
 The following wrappers are available and require no additional include
 directives:
 :cpp:class:`any`,
-:cpp:class:`bytes`, :cpp:class:`callable`, :cpp:class:`capsule`,
+:cpp:class:`bytearray`, :cpp:class:`bytes`, :cpp:class:`callable`, :cpp:class:`capsule`,
 :cpp:class:`dict`, :cpp:class:`ellipsis`, :cpp:class:`handle`,
 :cpp:class:`handle_t\<T\> <handle_t>`,
 :cpp:class:`bool_`, :cpp:class:`int_`, :cpp:class:`float_`,

include/nanobind/nb_lib.h

@@ -134,6 +134,14 @@ NB_CORE PyObject *bytes_from_cstr_and_size(const void *c, size_t n);
 
 // ========================================================================
 
+/// Convert a Python object into a Python byte array
+NB_CORE PyObject *bytearray_from_obj(PyObject *o);
+
+/// Convert a memory region into a Python byte array
+NB_CORE PyObject *bytearray_from_cstr_and_size(const void *c, size_t n);
+
+// ========================================================================
+
 /// Convert a Python object into a Python boolean object
 NB_CORE PyObject *bool_from_obj(PyObject *o);
 

include/nanobind/nb_types.h

@@ -443,6 +443,35 @@ class bytes : public object {
     size_t size() const { return (size_t) PyBytes_Size(m_ptr); }
 };
 
+class bytearray : public object {
+    NB_OBJECT(bytearray, object, "bytearray", PyByteArray_Check)
+
+#if PY_VERSION_HEX >= 0x03090000
+    bytearray()
+        : object(PyObject_CallNoArgs((PyObject *)&PyByteArray_Type), detail::steal_t{}) { }
+#else
+    bytearray()
+        : object(PyObject_CallObject((PyObject *)&PyByteArray_Type, NULL), detail::steal_t{}) { }
+#endif
+
+    explicit bytearray(handle h)
+        : object(detail::bytearray_from_obj(h.ptr()), detail::steal_t{}) { }
+
+    explicit bytearray(const void *s, size_t n)
+        : object(detail::bytearray_from_cstr_and_size(s, n), detail::steal_t{}) { }
+
+    const char *c_str() const { return PyByteArray_AsString(m_ptr); }
+
+    const void *data() const { return (const void *) PyByteArray_AsString(m_ptr); }
+
+    size_t size() const { return (size_t) PyByteArray_Size(m_ptr); }
+
+    void resize(size_t n) {
+        if (PyByteArray_Resize(m_ptr, (Py_ssize_t) n) != 0)
+            detail::raise_python_error();
+    }
+};
+
 class tuple : public object {
     NB_OBJECT(tuple, object, "tuple", PyTuple_Check)
     tuple() : object(PyTuple_New(0), detail::steal_t()) { }

src/common.cpp

@@ -579,6 +579,23 @@ PyObject *bytes_from_cstr_and_size(const void *str, size_t size) {
     return result;
 }
 
+// ========================================================================
+
+PyObject *bytearray_from_obj(PyObject *o) {
+    PyObject *result = PyByteArray_FromObject(o);
+    if (!result)
+        raise_python_error();
+    return result;
+}
+
+PyObject *bytearray_from_cstr_and_size(const void *str, size_t size) {
+    PyObject *result = PyByteArray_FromStringAndSize((const char *) str, (Py_ssize_t) size);
+    if (!result)
+        raise_python_error();
+    return result;
+}
+
+
 // ========================================================================
 
 PyObject *bool_from_obj(PyObject *o) {

tests/test_functions.cpp

@@ -361,4 +361,12 @@ NB_MODULE(test_functions_ext, m) {
     });
 
     m.def("hash_it", [](nb::handle h) { return nb::hash(h); });
+
+    // Test bytearray type
+    m.def("test_bytearray_new",     []() { return nb::bytearray(); });
+    m.def("test_bytearray_new",     [](const char *c, int size) { return nb::bytearray(c, size); });
+    m.def("test_bytearray_copy",    [](nb::bytearray o) { return nb::bytearray(o.c_str(), o.size()); });
+    m.def("test_bytearray_c_str",   [](nb::bytearray o) -> const char * { return o.c_str(); });
+    m.def("test_bytearray_size",    [](nb::bytearray o) { return o.size(); });
+    m.def("test_bytearray_resize",  [](nb::bytearray c, int size) { return c.resize(size); });
 }

tests/test_functions.py

@@ -598,3 +598,36 @@ def test43_wrappers_set():
 def test44_hash():
     value = (1, 2, 3)
     assert t.hash_it(value) == hash(value);
+
+
+def test45_new():
+    assert t.test_bytearray_new() == bytearray()
+    assert t.test_bytearray_new("\x00\x01\x02\x03", 4) == bytearray(
+        b"\x00\x01\x02\x03"
+    )
+    assert t.test_bytearray_new("", 0) == bytearray()
+
+
+def test46_copy():
+    o = bytearray(b"\x00\x01\x02\x03")
+    c = t.test_bytearray_copy(o)
+    assert c == o
+    o.clear()
+    assert c != o
+
+
+def test47_c_str():
+    o = bytearray(b"Hello, world!")
+    assert t.test_bytearray_c_str(o) == "Hello, world!"
+
+
+def test48_size():
+    o = bytearray(b"Hello, world!")
+    assert t.test_bytearray_size(o) == len(o)
+
+
+def test49_resize():
+    o = bytearray(b"\x00\x01\x02\x03")
+    assert len(o) == 4
+    t.test_bytearray_resize(o, 8)
+    assert len(o) == 8

tests/test_functions_ext.pyi.ref

@@ -140,6 +140,20 @@ def test_args_kwonly_kwargs(i: int, j: float, *args, z: int, **kwargs) -> tuple:
 
 def test_bad_tuple() -> tuple: ...
 
+def test_bytearray_c_str(arg: bytearray, /) -> str: ...
+
+def test_bytearray_copy(arg: bytearray, /) -> bytearray: ...
+
+@overload
+def test_bytearray_new() -> bytearray: ...
+
+@overload
+def test_bytearray_new(arg0: str, arg1: int, /) -> bytearray: ...
+
+def test_bytearray_resize(arg0: bytearray, arg1: int, /) -> None: ...
+
+def test_bytearray_size(arg: bytearray, /) -> int: ...
+
 def test_call_1(arg: object, /) -> object: ...
 
 def test_call_2(arg: object, /) -> object: ...