From 8e5fd142d9791f3f7ff66485de479578e64d9f37 Mon Sep 17 00:00:00 2001
From: Noah Kim <noahbkim@gmail.com>
Date: Tue, 13 Aug 2024 02:51:16 -0400
Subject: [PATCH] Add bytearray wrapper type (#654)

* Add bytearray wrapper type
---
 docs/api_core.rst                | 38 +++++++++++++++++++++++++++++++-
 docs/changelog.rst               |  9 +++++++-
 docs/exchanging.rst              |  2 +-
 include/nanobind/nb_lib.h        |  8 +++++++
 include/nanobind/nb_types.h      | 29 ++++++++++++++++++++++++
 src/common.cpp                   | 17 ++++++++++++++
 tests/test_functions.cpp         |  8 +++++++
 tests/test_functions.py          | 33 +++++++++++++++++++++++++++
 tests/test_functions_ext.pyi.ref | 14 ++++++++++++
 9 files changed, 155 insertions(+), 3 deletions(-)

diff --git a/docs/api_core.rst b/docs/api_core.rst
index 3976b1d2..b90191a4 100644
--- a/docs/api_core.rst
+++ b/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.
diff --git a/docs/changelog.rst b/docs/changelog.rst
index da873a48..d960b494 100644
--- a/docs/changelog.rst
+++ b/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
diff --git a/docs/exchanging.rst b/docs/exchanging.rst
index d83eb753..148f510f 100644
--- a/docs/exchanging.rst
+++ b/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_`,
diff --git a/include/nanobind/nb_lib.h b/include/nanobind/nb_lib.h
index 5ae08a16..a1930917 100644
--- a/include/nanobind/nb_lib.h
+++ b/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);
 
diff --git a/include/nanobind/nb_types.h b/include/nanobind/nb_types.h
index 8c71e4e7..68a16a2a 100644
--- a/include/nanobind/nb_types.h
+++ b/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()) { }
diff --git a/src/common.cpp b/src/common.cpp
index 30f02817..dceeee5e 100644
--- a/src/common.cpp
+++ b/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) {
diff --git a/tests/test_functions.cpp b/tests/test_functions.cpp
index d11ad7e3..cee371cd 100644
--- a/tests/test_functions.cpp
+++ b/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); });
 }
diff --git a/tests/test_functions.py b/tests/test_functions.py
index 357da67a..2217e7d5 100644
--- a/tests/test_functions.py
+++ b/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
diff --git a/tests/test_functions_ext.pyi.ref b/tests/test_functions_ext.pyi.ref
index 45a6701d..dcaaca0c 100644
--- a/tests/test_functions_ext.pyi.ref
+++ b/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: ...