From 6098b26f3ece2096e14695fd02b040bd0658d2ac Mon Sep 17 00:00:00 2001
From: William S Fulton <wsf@fultondesigns.co.uk>
Date: Mon, 3 Apr 2023 22:09:39 +0100
Subject: Iterator Protocol support for std::array wrappers

Introduce swig::IteratorProtocol class and assign which can be partially
specialized by different container types, such as std::array.
---
 Lib/python/pycontainer.swg | 46 +++++++++++++-------------------------------
 Lib/python/std_array.i     | 48 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 33 deletions(-)

(limited to 'Lib')

diff --git a/Lib/python/pycontainer.swg b/Lib/python/pycontainer.swg
index a1affa86b..65223b615 100644
--- a/Lib/python/pycontainer.swg
+++ b/Lib/python/pycontainer.swg
@@ -1014,7 +1014,8 @@ namespace swig {
   }
 
   template <class Seq, class T = typename Seq::value_type >
-    static int assign_iterator_protocol(PyObject *obj, Seq **seq) {
+  struct IteratorProtocol {
+    static int assign(PyObject *obj, Seq **seq) {
       int ret = SWIG_ERROR;
       PyObject *iter = PyObject_GetIter(obj);
       if (iter) {
@@ -1046,12 +1047,21 @@ namespace swig {
 
       return ret;
     }
+  };
 
   template <class Seq, class T = typename Seq::value_type >
   struct traits_asptr_stdseq {
     typedef Seq sequence;
     typedef T value_type;
 
+    static bool is_iterable(PyObject *obj) {
+      PyObject *iter = PyObject_GetIter(obj);
+      bool is_iter = iter != 0;
+      Py_XDECREF(iter);
+      PyErr_Clear();
+      return is_iter;
+    }
+
     static int asptr(PyObject *obj, sequence **seq) {
       int ret = SWIG_ERROR;
       if (obj == Py_None || SWIG_Python_GetSwigThis(obj)) {
@@ -1061,6 +1071,8 @@ namespace swig {
 	  if (seq) *seq = p;
 	  return SWIG_OLDOBJ;
 	}
+      } else if (is_iterable(obj)) {
+        ret = IteratorProtocol<Seq, T>::assign(obj, seq);
       } else if (PySequence_Check(obj)) {
 	try {
 	  SwigPySequence_Cont<value_type> swigpyseq(obj);
@@ -1080,38 +1092,6 @@ namespace swig {
 	  }
 	  return SWIG_ERROR;
 	}
-      } else {
-#if 0
-        PyObject *iter = PyObject_GetIter(obj);
-        if (iter) {
-          PyObject *item = PyIter_Next(iter);
-          ret = SWIG_OK;
-	  if (seq)
-            *seq = new sequence();
-          while (item) {
-            try {
-              if (seq) {
-                (*seq)->insert((*seq)->end(), swig::as<value_type>(item));
-              } else {
-                if (!swig::check<value_type>(item))
-                  ret = SWIG_ERROR;
-              }
-            } catch (std::exception& e) {
-              if (seq) {
-                if (!PyErr_Occurred()) {
-                  PyErr_SetString(PyExc_TypeError, e.what());
-                }
-              }
-              ret = SWIG_ERROR;
-	    }
-            Py_DECREF(item);
-            item = (ret == SWIG_OK) ? PyIter_Next(iter) : 0;
-          }
-          Py_DECREF(iter);
-        }
-#else
-      ret = swig::assign_iterator_protocol<Seq, T>(obj, seq);
-#endif
       }
       return ret;
     }
diff --git a/Lib/python/std_array.i b/Lib/python/std_array.i
index a3de3125b..9cca563c0 100644
--- a/Lib/python/std_array.i
+++ b/Lib/python/std_array.i
@@ -29,6 +29,54 @@
       std::copy(swigpyseq.begin(), swigpyseq.end(), seq->begin());
     }
 
+    template <class T, size_t N>
+    struct IteratorProtocol<std::array<T, N>, T> {
+      static int assign(PyObject *obj, std::array<T, N> **seq) {
+        int ret = SWIG_ERROR;
+        PyObject *iter = PyObject_GetIter(obj);
+        if (iter) {
+          PyObject *item = PyIter_Next(iter);
+          size_t count = 0;
+          typename std::array<T, N>::iterator array_iter = nullptr;
+          ret = SWIG_OK;
+          if (seq) {
+            *seq = new std::array<T, N>();
+            array_iter = (*seq)->begin();
+          }
+          while (item && (count < N)) {
+            try {
+              if (seq) {
+                *array_iter++ = swig::as<T>(item);
+              } else {
+                if (!swig::check<T>(item))
+                  ret = SWIG_ERROR;
+              }
+            } catch (std::exception& e) {
+              if (seq) {
+                if (!PyErr_Occurred()) {
+                  PyErr_SetString(PyExc_TypeError, e.what());
+                }
+              }
+              ret = SWIG_ERROR;
+            }
+            ++count;
+            Py_DECREF(item);
+            item = (ret == SWIG_OK) ? PyIter_Next(iter) : 0;
+          }
+          if ((ret == SWIG_OK) && (count != N || item)) {
+            PyErr_SetString(PyExc_TypeError, "std::array size does not match source container size");
+            ret = SWIG_ERROR;
+          }
+          Py_XDECREF(item);
+          Py_DECREF(iter);
+          if (seq && (ret == SWIG_ERROR))
+            delete *seq;
+        }
+
+        return ret;
+      }
+    };
+
     template <class T, size_t N>
     inline void
     erase(std::array<T, N>* SWIGUNUSEDPARM(seq), const typename std::array<T, N>::iterator& SWIGUNUSEDPARM(position)) {
-- 
cgit v1.2.1