From 417067eeb8d93e8ed49f3ff580712ae2b9a85259 Mon Sep 17 00:00:00 2001
From: Vikram Pal <vikrampal659@gmail.com>
Date: Sun, 14 Feb 2021 20:21:13 +0530
Subject: [PATCH] Add pybind11::bytearray (#2799)

* Add initial implementation

* Add few more methods

* Add tests

* Fix a typo

* Use std::string constructor which takes size

* Fix implicit sign conversion error

* Add size method and test

* Remove implicit conversion

* Fix bytearray constructors and operator std::string()

* Make implicit bytearray constructor explicit

* Rerun tests

* Add null check

* Rerun tests

* Rerun tests - 2

* Remove NULL check
---
 include/pybind11/pytypes.h | 28 ++++++++++++++++++++++++++++
 tests/test_pytypes.cpp     |  7 +++++++
 tests/test_pytypes.py      |  8 +++++++-
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/include/pybind11/pytypes.h b/include/pybind11/pytypes.h
index 23a09bd8b..6f0e6067e 100644
--- a/include/pybind11/pytypes.h
+++ b/include/pybind11/pytypes.h
@@ -1070,6 +1070,34 @@ inline str::str(const bytes& b) {
     m_ptr = obj.release().ptr();
 }
 
+/// \addtogroup pytypes
+/// @{
+class bytearray : public object {
+public:
+    PYBIND11_OBJECT_CVT(bytearray, object, PyByteArray_Check, PyByteArray_FromObject)
+
+    bytearray(const char *c, size_t n)
+        : object(PyByteArray_FromStringAndSize(c, (ssize_t) n), stolen_t{}) {
+        if (!m_ptr) pybind11_fail("Could not allocate bytearray object!");
+    }
+
+    bytearray()
+        : bytearray("", 0) {}
+
+    explicit bytearray(const std::string &s) : bytearray(s.data(), s.size()) { }
+
+    size_t size() const { return static_cast<size_t>(PyByteArray_Size(m_ptr)); }
+
+    explicit operator std::string() const {
+        char *buffer = PyByteArray_AS_STRING(m_ptr);
+        ssize_t size = PyByteArray_GET_SIZE(m_ptr);
+        return std::string(buffer, static_cast<size_t>(size));
+    }
+};
+// Note: breathe >= 4.17.0 will fail to build docs if the below two constructors
+// are included in the doxygen group; close here and reopen after as a workaround
+/// @} pytypes
+
 /// \addtogroup pytypes
 /// @{
 class none : public object {
diff --git a/tests/test_pytypes.cpp b/tests/test_pytypes.cpp
index f07b1fba2..6921796aa 100644
--- a/tests/test_pytypes.cpp
+++ b/tests/test_pytypes.cpp
@@ -92,6 +92,10 @@ TEST_SUBMODULE(pytypes, m) {
     m.def("bytes_from_string", []() { return py::bytes(std::string("foo")); });
     m.def("bytes_from_str", []() { return py::bytes(py::str("bar", 3)); });
 
+    // test bytearray
+    m.def("bytearray_from_string", []() { return py::bytearray(std::string("foo")); });
+    m.def("bytearray_size", []() { return py::bytearray("foo").size(); });
+
     // test_capsule
     m.def("return_capsule_with_destructor", []() {
         py::print("creating capsule");
@@ -210,6 +214,7 @@ TEST_SUBMODULE(pytypes, m) {
     m.def("default_constructors", []() {
         return py::dict(
             "bytes"_a=py::bytes(),
+            "bytearray"_a=py::bytearray(),
             "str"_a=py::str(),
             "bool"_a=py::bool_(),
             "int"_a=py::int_(),
@@ -224,6 +229,7 @@ TEST_SUBMODULE(pytypes, m) {
     m.def("converting_constructors", [](py::dict d) {
         return py::dict(
             "bytes"_a=py::bytes(d["bytes"]),
+            "bytearray"_a=py::bytearray(d["bytearray"]),
             "str"_a=py::str(d["str"]),
             "bool"_a=py::bool_(d["bool"]),
             "int"_a=py::int_(d["int"]),
@@ -240,6 +246,7 @@ TEST_SUBMODULE(pytypes, m) {
         // When converting between Python types, obj.cast<T>() should be the same as T(obj)
         return py::dict(
             "bytes"_a=d["bytes"].cast<py::bytes>(),
+            "bytearray"_a=d["bytearray"].cast<py::bytearray>(),
             "str"_a=d["str"].cast<py::str>(),
             "bool"_a=d["bool"].cast<py::bool_>(),
             "int"_a=d["int"].cast<py::int_>(),
diff --git a/tests/test_pytypes.py b/tests/test_pytypes.py
index cebbd4791..f17bf76b7 100644
--- a/tests/test_pytypes.py
+++ b/tests/test_pytypes.py
@@ -143,6 +143,11 @@ def test_bytes(doc):
     )
 
 
+def test_bytearray(doc):
+    assert m.bytearray_from_string().decode() == "foo"
+    assert m.bytearray_size() == len("foo")
+
+
 def test_capsule(capture):
     pytest.gc_collect()
     with capture:
@@ -223,7 +228,7 @@ def test_accessors():
 
 def test_constructors():
     """C++ default and converting constructors are equivalent to type calls in Python"""
-    types = [bytes, str, bool, int, float, tuple, list, dict, set]
+    types = [bytes, bytearray, str, bool, int, float, tuple, list, dict, set]
     expected = {t.__name__: t() for t in types}
     if env.PY2:
         # Note that bytes.__name__ == 'str' in Python 2.
@@ -234,6 +239,7 @@ def test_constructors():
 
     data = {
         bytes: b"41",  # Currently no supported or working conversions.
+        bytearray: bytearray(b"41"),
         str: 42,
         bool: "Not empty",
         int: "42",