Search code examples
pythonc++iteratorswig

How to make a C++ class iterable from Python using SWIG?


I have a C++ class Collection that manages a std::vector<Element> (a private member of the class).

From C++ I can iterate through the vector using the begin() and end() iterators (which are just typedefs for the vector's iterators) like:

Collection col;
for (Collection::const_iterator itr = col.begin(); itr != col.end(); itr++)
{
  std::cout << itr->get() << std::endl;
}

Now I wish to do a similar thing from Python like:

import example
el = example.Element()
el.set(5)
col = example.Collection()
col.add(el)
for e in col:
    print e.get()

But this results in:

TypeError: 'Collection' object is not iterable

I am not able to configure SWIG in a way that it generates the __iter__ (I think it's the only thing it needs) for the Python Collection class. How am I supposed to do this?

This is my code:

example.h:

#include <vector>

class Element
{
public:
  Element();
  ~Element();

  int get() const;
  void set(const int var);

private:
  int variable_;
};

class Collection
{
public:
  Collection();
  ~Collection();

  void add(const Element& element);

  typedef std::vector<Element> tElements;

  // iterators
  typedef tElements::iterator iterator;
  typedef tElements::const_iterator const_iterator;
  iterator begin();
  const_iterator begin() const;
  iterator end();
  const_iterator end() const;

private:
  tElements          elements_;
};

example.cpp:

#include "example.h"

Element::Element() {}

Element::~Element() {}

int Element::get() const
{
  return variable_;
}

void Element::set(const int var)
{
  variable_ = var;
}

Collection::Collection() : elements_() {}

Collection::~Collection() {}

void Collection::add(const Element& element)
{
  elements_.push_back(element);
}

Collection::iterator Collection::begin()
{
  return elements_.begin();
}

Collection::const_iterator Collection::begin() const
{
  return elements_.begin();
}

Collection::iterator Collection::end()
{
  return elements_.end();
}

Collection::const_iterator Collection::end() const
{
  return elements_.end();
}

example.i:

%module example
%{
#include "example.h"
%}

// I've tried to add this, but that generates a whole
// other class, that is not what I want.
// %include "std_vector.i"
// %template(ElementVector) std::vector<Element>;

// I've also tried to %extend the class (which I think is what I want,
// but I cannot figure out with what to extend it with)

// Include the header file with above prototypes
%include "example.h"

Compile with:

swig -python -c++ -o example_wrap.cpp example.i
g++ -fPIC -c example.cpp example_wrap.cpp -I/usr/include/python2.6
g++ -shared example.o example_wrap.o -o _example.so

Solution

  • Inspired by the last example of: https://stackoverflow.com/a/8828454/3613373. I came up with a slightly different approach that does not use a variable to check the StopIterator exception state.

    Also it only uses the begin() end end() iterators of Collection without requiring to expose (make public) the std::vector<Element> itself.

    example.i:

    %module example
    %{
    #include "example.h"
    %}
    
    %inline %{
    class StopIterator {};
    class Iterator {
      public:
        Iterator(Collection::iterator _cur, Collection::iterator _end) : cur(_cur), end(_end) {}
        Iterator* __iter__()
        {
          return this;
        }
        Collection::iterator cur;
        Collection::iterator end;
      };
    %}
    
    %include "example.h"
    
    %include "exception.i"
    %exception Iterator::next {
      try
      {
        $action // calls %extend function next() below
      }
      catch (StopIterator)
      {
        PyErr_SetString(PyExc_StopIteration, "End of iterator");
        return NULL;
      }
    }
    
    %extend Iterator
    {
      Element& next()
      {
        if ($self->cur != $self->end)
        {
          // dereference the iterator and return reference to the object,
          // after that it increments the iterator
          return *$self->cur++;
        }
        throw StopIterator();
      }
    }
    
    %extend Collection {
      Iterator __iter__()
      {
        // return a constructed Iterator object
        return Iterator($self->begin(), $self->end());
      }
    };
    

    Only thing I could not figure out was how to create a templated version of Iterator in a way that I can pass just any Iterator<Collection, Element> without having to redefine next() for each template instantation. Solutions are welcome ;)