Change list_caster to also accept generator objects (PyGen_Check(src.ptr()).

rwgk · rwgk · commit 7d30378bdffa · 2023-05-30T11:59:16.000-07:00
Note for completeness: This is a more conservative change than google/pybind11clif#30042
diff --git a/include/pybind11/stl.h b/include/pybind11/stl.h
@@ -166,13 +166,38 @@ struct list_caster {
     using value_conv = make_caster<Value>;
 
     bool load(handle src, bool convert) {
-        if (!isinstance<sequence>(src) || isinstance<bytes>(src) || isinstance<str>(src)) {
+        if (isinstance<bytes>(src) || isinstance<str>(src)) {
             return false;
         }
-        auto s = reinterpret_borrow<sequence>(src);
+        if (isinstance<sequence>(src)) {
+            return convert_elements(src, convert);
+        }
+        if (!convert) {
+            return false;
+        }
+        if (PyGen_Check(src.ptr())) {
+            // Designed to be behavior-equivalent to passing tuple(src) from Python:
+            // The conversion to a tuple will first exhaust the generator object, to ensure that
+            // the generator is not left in an unpredictable (to the caller) partially-consumed
+            // state.
+            assert(isinstance<iterable>(src));
+            return convert_elements(tuple(reinterpret_borrow<iterable>(src)), convert);
+        }
+        return false;
+    }
+
+private:
+    template <typename T = Type, enable_if_t<has_reserve_method<T>::value, int> = 0>
+    void reserve_maybe(const sequence &s, Type *) {
+        value.reserve(s.size());
+    }
+    void reserve_maybe(const sequence &, void *) {}
+
+    bool convert_elements(handle seq, bool convert) {
+        auto s = reinterpret_borrow<sequence>(seq);
         value.clear();
         reserve_maybe(s, &value);
-        for (auto it : s) {
+        for (auto it : seq) {
             value_conv conv;
             if (!conv.load(it, convert)) {
                 return false;
@@ -182,13 +207,6 @@ struct list_caster {
         return true;
     }
 
-private:
-    template <typename T = Type, enable_if_t<has_reserve_method<T>::value, int> = 0>
-    void reserve_maybe(const sequence &s, Type *) {
-        value.reserve(s.size());
-    }
-    void reserve_maybe(const sequence &, void *) {}
-
 public:
     template <typename T>
     static handle cast(T &&src, return_value_policy policy, handle parent) {
diff --git a/tests/test_stl.py b/tests/test_stl.py
@@ -385,14 +385,13 @@ def test_pass_std_vector_int():
     fn = m.pass_std_vector_int
     assert fn([1, 2]) == 2
     assert fn((1, 2)) == 2
+    assert fn(i for i in range(3)) == 3
     with pytest.raises(TypeError):
         fn(set())
     with pytest.raises(TypeError):
         fn({})
     with pytest.raises(TypeError):
         fn({}.keys())
-    with pytest.raises(TypeError):
-        fn(i for i in range(3))
 
 
 def test_pass_std_set_int():
@@ -408,3 +407,13 @@ def test_pass_std_set_int():
         fn({}.keys())
     with pytest.raises(TypeError):
         fn(i for i in range(3))
+
+
+def test_list_caster_fully_consumes_generator_object():
+    def gen_mix():
+        yield from [1, 2.0, 3]
+
+    gen_obj = gen_mix()
+    with pytest.raises(TypeError):
+        m.pass_std_vector_int(gen_obj)
+    assert not tuple(gen_obj)
