Add bytearray wrapper type (#654)

* Add bytearray wrapper type

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

@@ -42,6 +42,8 @@ Version 2.1.0 (Aug 11, 2024)
   :cpp:struct:`nb::new_ <new_>`, passed in the same way they would be with
   :cpp:struct:`nb::init <init>`. (issue `#668
   <https://github.com/wjakob/nanobind/issues/668>`__)
+  
+* Added a wrapper type :cpp:class:`bytearray`.
 
 * Ability to use :cpp:func:`nb::cast <cast>` to create object with the
   :cpp:enumerator:`nb::rv_policy::reference_internal

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: ...