Commit 7602a0ef authored by Hastings Greer's avatar Hastings Greer
Browse files

Make vtkCollection implement iterable in the python wrapping

parent 04422983
Pipeline #16480 passed with stage
......@@ -9,6 +9,7 @@ vtk_add_test_python(
TestEnums.py
TestExecuteMethodFinalizeCrash.py
TestGhost.py
TestIterateCollection.py
TestMutable.py
TestNumpySupport.py
TestNumpyInterface.py
......
"""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')])
......@@ -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"
......
......@@ -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");
}
}
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment