Make vtkCollection implement iterable in the python wrapping

Author: HastingsGreer

Common/Core/Testing/Python/CMakeLists.txt

@@ -9,6 +9,7 @@ vtk_add_test_python(
   TestEnums.py
   TestExecuteMethodFinalizeCrash.py
   TestGhost.py
+  TestIterateCollection.py
   TestMutable.py
   TestNumpySupport.py
   TestNumpyInterface.py

Common/Core/Testing/Python/TestIterateCollection.py

@@ -0,0 +1,47 @@
+"""Test iterating through a vtkCollection with the standard python syntax"""
+
+
+import vtk
+from vtk.test import Testing
+
+
+class TestIterateCollection(Testing.vtkTest):
+    def setUp(self):
+        self.vtkObjs = [vtk.vtkObject() for _ in range(30)]
+        self.collection = vtk.vtkCollection()
+        for obj in self.vtkObjs:
+            self.collection.AddItem(obj)
+
+        self.emptyCollection = vtk.vtkCollection()
+
+    def testIterateCollection(self):
+        newObjsList = [obj for obj in self.collection]
+        self.assertEqual(self.vtkObjs, newObjsList)
+
+        counter = 0
+        for _ in self.collection:
+            counter += 1
+        self.assertEqual(counter, 30)
+
+        counter = 0
+        for _ in self.emptyCollection:
+           counter += 1
+        self.assertEqual(counter, 0)
+
+    def testCollectionChild(self):
+        #check that iteration is being inherited correctly
+        dataArray = vtk.vtkIntArray()
+
+        dataArrayCollection = vtk.vtkDataArrayCollection()
+        dataArrayCollection.AddItem(dataArray)
+
+        self.assertEqual([obj for obj in dataArrayCollection],
+                         [dataArray])
+
+    def testOperators(self):
+        self.assertTrue(self.vtkObjs[0] in self.collection)
+        self.assertEqual(list(self.collection), self.vtkObjs)
+
+
+if __name__ == "__main__":
+    Testing.main([(TestIterateCollection, 'test')])

Wrapping/Tools/vtkWrapPythonClass.c

@@ -550,9 +550,29 @@ void vtkWrapPython_GenerateObjectType(
     "  PyVTKObject_Traverse, // tp_traverse\n"
     "  0, // tp_clear\n"
     "  0, // tp_richcompare\n"
-    "  offsetof(PyVTKObject, vtk_weakreflist), // tp_weaklistoffset\n"
-    "  0, // tp_iter\n"
-    "  0, // tp_iternext\n"
+    "  offsetof(PyVTKObject, vtk_weakreflist), // tp_weaklistoffset\n");
+  if (strcmp(classname, "vtkCollection") == 0)
+    {
+    fprintf(fp,
+      "  PyvtkCollection_iter, // tp_iter\n"
+      "  0, // tp_iternext\n");
+    }
+  else
+    {
+    if(strcmp(classname, "vtkCollectionIterator") == 0)
+      {
+      fprintf(fp,
+        "  PyvtkCollectionIterator_iter, // tp_iter\n"
+        "  PyvtkCollectionIterator_next, // tp_iternext\n");
+      }
+    else
+      {
+      fprintf(fp,
+        "  0, // tp_iter\n"
+        "  0, // tp_iternext\n");
+      }
+    }
+  fprintf(fp,
     "  0, // tp_methods\n"
     "  0, // tp_members\n"
     "  PyVTKObject_GetSet, // tp_getset\n"

Wrapping/Tools/vtkWrapPythonMethodDef.c

@@ -873,4 +873,58 @@ static void vtkWrapPython_CustomMethods(
             "\n",
             classname, data->Name, data->Name, data->Name);
     }
+    if (strcmp("vtkCollection", data->Name) == 0 &&
+       do_constructors == 0)
+     {
+      fprintf(fp,
+              "static PyObject *\n"
+              "PyvtkCollection_iter(PyObject *self)\n"
+              "{\n"
+              "  PyVTKObject* vp = (PyVTKObject *) self;\n"
+              "  vtkCollection *op = (vtkCollection* ) vp->vtk_ptr; \n"
+              "  \n"
+              "  PyObject *result = NULL;\n"
+              "  \n"
+              "  if (op)\n"
+              "    {\n"
+              "    vtkObject *tempr = (vtkObject *) op->NewIterator();\n"
+              "    if (tempr != NULL)\n"
+              "      {\n"
+              "      result = vtkPythonArgs::BuildVTKObject(tempr);\n"
+              "      }\n"
+              "    }\n"
+              "\n"
+              "  return result;\n"
+              "}\n");
+    }
+   if (strcmp("vtkCollectionIterator", data->Name) == 0 &&
+      do_constructors == 0)
+    {
+     fprintf(fp,
+              "static PyObject * PyvtkCollectionIterator_next(PyObject *self)\n"
+              "{\n"
+              "  PyVTKObject* vp = (PyVTKObject *) self;\n"
+              "  vtkCollectionIterator *op = (vtkCollectionIterator* ) vp->vtk_ptr; \n"
+              "  \n"
+              "  PyObject *result = NULL;\n"
+              "  \n"
+              "  if (op)\n"
+              "    {\n"
+              "    vtkObject *tempr = op->GetCurrentObject();\n"
+              "    op->GoToNextItem();\n"
+              "    if (tempr != NULL)\n"
+              "      {\n"
+              "      result = vtkPythonArgs::BuildVTKObject(tempr);\n"
+              "      }\n"
+              "    }\n"
+              "\n"
+              "  return result;\n"
+              "}\n"
+              "\n"
+              "static PyObject *PyvtkCollectionIterator_iter(PyObject *self)\n"
+              "{\n"
+              "  Py_INCREF(self);\n"
+              "  return self;\n"
+              "}\n");
+    }
 }