Add support for all STL containers to be constructible from a Python set

3 files changed, 61 insertions, 66 deletions

diff --git a/Lib/python/pyclasses.swg b/Lib/python/pyclasses.swg
index 31ebdd2a1..39c4e0316 100644
--- a/Lib/python/pyclasses.swg
+++ b/Lib/python/pyclasses.swg
@@ -11,7 +11,7 @@
   or as a member variable:
   
      struct A {
-       SwigPtr_PyObject obj;
+       SwigPtr_PyObject _obj;
        A(PyObject *o) : _obj(o) {
        }
      };
diff --git a/Lib/python/pycontainer.swg b/Lib/python/pycontainer.swg
index 65223b615..7012b0b1d 100644
--- a/Lib/python/pycontainer.swg
+++ b/Lib/python/pycontainer.swg
@@ -1015,36 +1015,28 @@ namespace swig {
 
   template <class Seq, class T = typename Seq::value_type >
   struct IteratorProtocol {
-    static int assign(PyObject *obj, Seq **seq) {
-      int ret = SWIG_ERROR;
-      PyObject *iter = PyObject_GetIter(obj);
+    static void assign(PyObject *obj, Seq *seq) {
+      SwigVar_PyObject iter = PyObject_GetIter(obj);
       if (iter) {
-        PyObject *item = PyIter_Next(iter);
-        ret = SWIG_OK;
-        if (seq)
-          *seq = new Seq();
+        SwigVar_PyObject item = PyIter_Next(iter);
         while (item) {
-          try {
-            if (seq) {
-              (*seq)->insert((*seq)->end(), 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;
-          }
-          Py_DECREF(item);
-          item = (ret == SWIG_OK) ? PyIter_Next(iter) : 0;
+          seq->insert(seq->end(), swig::as<T>(item));
+          item = PyIter_Next(iter);
         }
-        Py_DECREF(iter);
       }
+    }
 
+    static bool check(PyObject *obj) {
+      int ret = false;
+      SwigVar_PyObject iter = PyObject_GetIter(obj);
+      if (iter) {
+        SwigVar_PyObject item = PyIter_Next(iter);
+        ret = true;
+        while (item) {
+          ret = swig::check<T>(item);
+          item = ret ? PyIter_Next(iter) : 0;
+        }
+      }
       return ret;
     }
   };
@@ -1055,11 +1047,9 @@ namespace swig {
     typedef T value_type;
 
     static bool is_iterable(PyObject *obj) {
-      PyObject *iter = PyObject_GetIter(obj);
-      bool is_iter = iter != 0;
-      Py_XDECREF(iter);
+      SwigVar_PyObject iter = PyObject_GetIter(obj);
       PyErr_Clear();
-      return is_iter;
+      return iter != 0;
     }
 
     static int asptr(PyObject *obj, sequence **seq) {
@@ -1072,7 +1062,22 @@ namespace swig {
 	  return SWIG_OLDOBJ;
 	}
       } else if (is_iterable(obj)) {
-        ret = IteratorProtocol<Seq, T>::assign(obj, seq);
+	try {
+	  if (seq) {
+	    *seq = new sequence();
+            IteratorProtocol<Seq, T>::assign(obj, *seq);
+            if (!PyErr_Occurred())
+              return SWIG_NEWOBJ;
+	  } else {
+	    return IteratorProtocol<Seq, T>::check(obj) ? SWIG_OK : SWIG_ERROR;
+	  }
+	} catch (std::exception& e) {
+          if (seq && !PyErr_Occurred())
+            PyErr_SetString(PyExc_TypeError, e.what());
+	}
+        if (seq)
+          delete *seq;
+	return SWIG_ERROR;
       } else if (PySequence_Check(obj)) {
 	try {
 	  SwigPySequence_Cont<value_type> swigpyseq(obj);
diff --git a/Lib/python/std_array.i b/Lib/python/std_array.i
index 9cca563c0..d41d1f15e 100644
--- a/Lib/python/std_array.i
+++ b/Lib/python/std_array.i
@@ -31,48 +31,38 @@
 
     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);
+
+      static void assign(PyObject *obj, std::array<T, N> *seq) {
+        SwigVar_PyObject iter = PyObject_GetIter(obj);
         if (iter) {
-          PyObject *item = PyIter_Next(iter);
+          SwigVar_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();
-          }
+          typename std::array<T, N>::iterator 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;
+            *array_iter++ = swig::as<T>(item);
+            item = PyIter_Next(iter);
           }
-          Py_XDECREF(item);
-          Py_DECREF(iter);
-          if (seq && (ret == SWIG_ERROR))
-            delete *seq;
+          if (count != N || item)
+            throw std::invalid_argument("std::array size does not match source container size");
         }
+      }
 
+      static bool check(PyObject *obj) {
+        int ret = false;
+        SwigVar_PyObject iter = PyObject_GetIter(obj);
+        if (iter) {
+          SwigVar_PyObject item = PyIter_Next(iter);
+          size_t count = 0;
+          ret = true;
+          while (item && (count < N)) {
+            ++count;
+            ret = swig::check<T>(item);
+            item = ret ? PyIter_Next(iter) : 0;
+          }
+          if (count != N || item)
+            ret = false;
+        }
         return ret;
       }
     };