From ae185b7f1973aef0bd15f1f20cbcbac7d3f3abb5 Mon Sep 17 00:00:00 2001
From: Jason Rhinelander <jason@imaginary.ca>
Date: Wed, 7 Dec 2016 18:43:29 -0500
Subject: [PATCH] std::valarray support for stl.h (#545)

* Added ternary support with descr args

Current the `_<bool>(a, b)` ternary support only works for `char[]` `a`
and `b`; this commit allows it to work for `descr` `a` and `b` arguments
as well.

* Add support for std::valarray to stl.h

This abstracts the std::array into a `array_caster` which can then be
used with either std::array or std::valarray, the main difference being
that std::valarray is resizable.  (It also lets the array_caster be
potentially used for other std::array-like interfaces, much as the
list_caster and map_caster currently provide).

* Small stl.h cleanups

- Remove redundant `type` typedefs
- make internal list_caster methods private
---
 include/pybind11/descr.h    |  6 +++++
 include/pybind11/stl.h      | 44 +++++++++++++++++++++++++++----------
 tests/test_python_types.cpp | 12 ++++++++++
 tests/test_python_types.py  | 14 ++++++++++++
 4 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/include/pybind11/descr.h b/include/pybind11/descr.h
index 22e94ef9..2c3fb3d1 100644
--- a/include/pybind11/descr.h
+++ b/include/pybind11/descr.h
@@ -81,6 +81,10 @@ template <bool B, size_t Size1, size_t Size2>
 constexpr enable_if_t<!B, descr<Size2 - 1, 0>> _(char const(&)[Size1], char const(&text2)[Size2]) {
     return _(text2);
 }
+template <bool B, size_t SizeA1, size_t SizeA2, size_t SizeB1, size_t SizeB2>
+constexpr enable_if_t<B, descr<SizeA1, SizeA2>> _(descr<SizeA1, SizeA2> d, descr<SizeB1, SizeB2>) { return d; }
+template <bool B, size_t SizeA1, size_t SizeA2, size_t SizeB1, size_t SizeB2>
+constexpr enable_if_t<!B, descr<SizeB1, SizeB2>> _(descr<SizeA1, SizeA2>, descr<SizeB1, SizeB2> d) { return d; }
 
 template <size_t Size> auto constexpr _() -> decltype(int_to_str<Size / 10, Size % 10>::digits) {
     return int_to_str<Size / 10, Size % 10>::digits;
@@ -154,6 +158,8 @@ PYBIND11_NOINLINE inline descr _(const char *text) {
 
 template <bool B> PYBIND11_NOINLINE enable_if_t<B, descr> _(const char *text1, const char *) { return _(text1); }
 template <bool B> PYBIND11_NOINLINE enable_if_t<!B, descr> _(char const *, const char *text2) { return _(text2); }
+template <bool B> PYBIND11_NOINLINE enable_if_t<B, descr> _(descr d, descr) { return d; }
+template <bool B> PYBIND11_NOINLINE enable_if_t<!B, descr> _(descr, descr d) { return d; }
 
 template <typename Type> PYBIND11_NOINLINE descr _() {
     const std::type_info *types[2] = { &typeid(Type), nullptr };
diff --git a/include/pybind11/stl.h b/include/pybind11/stl.h
index 766a8f21..d4b0fc91 100644
--- a/include/pybind11/stl.h
+++ b/include/pybind11/stl.h
@@ -16,6 +16,7 @@
 #include <unordered_map>
 #include <iostream>
 #include <list>
+#include <valarray>
 
 #if defined(_MSC_VER)
 #pragma warning(push)
@@ -72,7 +73,6 @@ template <typename Type, typename Key> struct set_caster {
 };
 
 template <typename Type, typename Key, typename Value> struct map_caster {
-    using type = Type;
     using key_conv   = make_caster<Key>;
     using value_conv = make_caster<Value>;
 
@@ -92,7 +92,7 @@ template <typename Type, typename Key, typename Value> struct map_caster {
         return true;
     }
 
-    static handle cast(const type &src, return_value_policy policy, handle parent) {
+    static handle cast(const Type &src, return_value_policy policy, handle parent) {
         dict d;
         for (auto const &kv: src) {
             auto key = reinterpret_steal<object>(key_conv::cast(kv.first, policy, parent));
@@ -104,11 +104,10 @@ template <typename Type, typename Key, typename Value> struct map_caster {
         return d.release();
     }
 
-    PYBIND11_TYPE_CASTER(type, _("Dict[") + key_conv::name() + _(", ") + value_conv::name() + _("]"));
+    PYBIND11_TYPE_CASTER(Type, _("Dict[") + key_conv::name() + _(", ") + value_conv::name() + _("]"));
 };
 
 template <typename Type, typename Value> struct list_caster {
-    using type = Type;
     using value_conv = make_caster<Value>;
 
     bool load(handle src, bool convert) {
@@ -126,11 +125,13 @@ template <typename Type, typename Value> struct list_caster {
         return true;
     }
 
+private:
     template <typename T = Type,
               enable_if_t<std::is_same<decltype(std::declval<T>().reserve(0)), void>::value, int> = 0>
     void reserve_maybe(sequence s, Type *) { value.reserve(s.size()); }
     void reserve_maybe(sequence, void *) { }
 
+public:
     static handle cast(const Type &src, return_value_policy policy, handle parent) {
         list l(src.size());
         size_t index = 0;
@@ -152,28 +153,40 @@ template <typename Type, typename Alloc> struct type_caster<std::vector<Type, Al
 template <typename Type, typename Alloc> struct type_caster<std::list<Type, Alloc>>
  : list_caster<std::list<Type, Alloc>, Type> { };
 
-template <typename Type, size_t Size> struct type_caster<std::array<Type, Size>> {
-    using array_type = std::array<Type, Size>;
-    using value_conv = make_caster<Type>;
+template <typename ArrayType, typename Value, bool Resizable, size_t Size = 0> struct array_caster {
+    using value_conv = make_caster<Value>;
 
+private:
+    template <bool R = Resizable>
+    bool require_size(enable_if_t<R, size_t> size) {
+        if (value.size() != size)
+            value.resize(size);
+        return true;
+    }
+    template <bool R = Resizable>
+    bool require_size(enable_if_t<!R, size_t> size) {
+        return size == Size;
+    }
+
+public:
     bool load(handle src, bool convert) {
         if (!isinstance<list>(src))
             return false;
         auto l = reinterpret_borrow<list>(src);
-        if (l.size() != Size)
+        if (!require_size(l.size()))
             return false;
         value_conv conv;
         size_t ctr = 0;
         for (auto it : l) {
             if (!conv.load(it, convert))
                 return false;
-            value[ctr++] = cast_op<Type>(conv);
+            value[ctr++] = cast_op<Value>(conv);
         }
         return true;
     }
 
-    static handle cast(const array_type &src, return_value_policy policy, handle parent) {
-        list l(Size);
+    static handle cast(const ArrayType &src, return_value_policy policy, handle parent) {
+        list l(src.size());
         size_t index = 0;
         for (auto const &value: src) {
             auto value_ = reinterpret_steal<object>(value_conv::cast(value, policy, parent));
@@ -183,9 +196,16 @@ template <typename Type, size_t Size> struct type_caster<std::array<Type, Size>>
         }
         return l.release();
     }
-    PYBIND11_TYPE_CASTER(array_type, _("List[") + value_conv::name() + _("[") + _<Size>() + _("]]"));
+
+    PYBIND11_TYPE_CASTER(ArrayType, _("List[") + value_conv::name() + _<Resizable>(_(""), _("[") + _<Size>() + _("]")) + _("]"));
 };
 
+template <typename Type, size_t Size> struct type_caster<std::array<Type, Size>>
+ : array_caster<std::array<Type, Size>, Type, false, Size> { };
+
+template <typename Type> struct type_caster<std::valarray<Type>>
+ : array_caster<std::valarray<Type>, Type, true> { };
+
 template <typename Key, typename Compare, typename Alloc> struct type_caster<std::set<Key, Compare, Alloc>>
   : set_caster<std::set<Key, Compare, Alloc>, Key> { };
 
diff --git a/tests/test_python_types.cpp b/tests/test_python_types.cpp
index 30fa3c79..68e07adb 100644
--- a/tests/test_python_types.cpp
+++ b/tests/test_python_types.cpp
@@ -77,6 +77,10 @@ public:
         return std::array<std::string, 2> {{ "array entry 1" , "array entry 2"}};
     }
 
+    std::valarray<int> get_valarray() {
+        return std::valarray<int>({ 1, 4, 9 });
+    }
+
     /* Easily iterate over a dictionary using a C++11 range-based for loop */
     void print_dict(py::dict dict) {
         for (auto item : dict)
@@ -132,6 +136,12 @@ public:
             py::print("array item {}: {}"_s.format(index++, item));
     }
 
+    void print_valarray(std::valarray<int> &varray) {
+        int index = 0;
+        for (auto item : varray)
+            py::print("valarray item {}: {}"_s.format(index++, item));
+    }
+
     void throw_exception() {
         throw std::runtime_error("This exception was intentionally thrown.");
     }
@@ -182,6 +192,7 @@ test_initializer python_types([](py::module &m) {
         .def("get_set", &ExamplePythonTypes::get_set, "Return a Python set")
         .def("get_set2", &ExamplePythonTypes::get_set_2, "Return a C++ set")
         .def("get_array", &ExamplePythonTypes::get_array, "Return a C++ array")
+        .def("get_valarray", &ExamplePythonTypes::get_valarray, "Return a C++ valarray")
         .def("print_dict", &ExamplePythonTypes::print_dict, "Print entries of a Python dictionary")
         .def("print_dict_2", &ExamplePythonTypes::print_dict_2, "Print entries of a C++ dictionary")
         .def("print_set", &ExamplePythonTypes::print_set, "Print entries of a Python set")
@@ -189,6 +200,7 @@ test_initializer python_types([](py::module &m) {
         .def("print_list", &ExamplePythonTypes::print_list, "Print entries of a Python list")
         .def("print_list_2", &ExamplePythonTypes::print_list_2, "Print entries of a C++ list")
         .def("print_array", &ExamplePythonTypes::print_array, "Print entries of a C++ array")
+        .def("print_valarray", &ExamplePythonTypes::print_valarray, "Print entries of a C++ valarray")
         .def("pair_passthrough", &ExamplePythonTypes::pair_passthrough, "Return a pair in reversed order")
         .def("tuple_passthrough", &ExamplePythonTypes::tuple_passthrough, "Return a triple in reversed order")
         .def("throw_exception", &ExamplePythonTypes::throw_exception, "Throw an exception")
diff --git a/tests/test_python_types.py b/tests/test_python_types.py
index abffa1a9..50a707ef 100644
--- a/tests/test_python_types.py
+++ b/tests/test_python_types.py
@@ -87,6 +87,15 @@ def test_instance(capture):
         array item 0: array entry 1
         array item 1: array entry 2
     """
+    varray_result = instance.get_valarray()
+    assert varray_result == [1, 4, 9]
+    with capture:
+        instance.print_valarray(varray_result)
+    assert capture.unordered == """
+        valarray item 0: 1
+        valarray item 1: 4
+        valarray item 2: 9
+    """
     with pytest.raises(RuntimeError) as excinfo:
         instance.throw_exception()
     assert str(excinfo.value) == "This exception was intentionally thrown."
@@ -164,6 +173,11 @@ def test_docs(doc):
 
         Return a C++ array
     """
+    assert doc(ExamplePythonTypes.get_valarray) == """
+        get_valarray(self: m.ExamplePythonTypes) -> List[int]
+
+        Return a C++ valarray
+    """
     assert doc(ExamplePythonTypes.print_dict) == """
         print_dict(self: m.ExamplePythonTypes, arg0: dict) -> None