Boost.CoroutineとBoost.OptionalとBoost.Pythonを使って、C++クラスの内部イテレータをPythonの外部イテレータとして見せる

python boost c++ coroutine

追記: うっかり戻り値ポリシーのところで boost::python::return_internal_reference<> を使ってしまい、メモリリークを引き起こしていたので修正しました。

Stackless Python であれば continuation が使えるので、Python 側で C++ をまたぐ外部イテレータと内部イテレータの反転が可能なんだけど、通常の Python ではやはり Boost.Coroutine に頼るほかない。

以下のサンプルでは、よくある Visitor パタンに従って実装されたクラスを想定していて、Visitor が visit されるたびに、Python 側に処理が戻るというフローを実現している。

container クラスの戻り値は新たに new されたオブジェクトなので、boost::python::return_value_policy を使いたいところだが、generator_iterator<> は、container クラスの持つ internal_iterator のインスタンスを弱参照しているため、generator_iterator<> のインスタンスが生きている限りは container クラスのインスタンスが生きていることを保証する必要がある。そこで、boost::python::with_custodian_and_ward_postcall<> を使った。テンプレートパラメータのところで owner_arg が 1 となっているのは、0 が戻り値 (lvalue) を表しているからだ。

もうひとつ特筆すべき点は、StopIterator を raise するところだ。Python の API を使った操作の結果 Python の例外が発生した場合は、boost::python::throw_error_already_set(); を呼ぶ。すると、内部で 大域脱出のための C++ の例外が投げられ、即座に Python 側に処理を戻すことができる。

#include <typeinfo>
#include <boost/python.hpp>
#include <boost/coroutine/generator.hpp>
#include <boost/optional.hpp>

// 何らかの visitor パタンを実装したクラス
class internal_iterator
{
public:
    internal_iterator(int start, int end): start_(start), end_(end) {}

    template<typename Thdlr_>
    void accept(Thdlr_& hdlr) const {
        // 内部イテレータ
        for (int i = start_; i < end_; ++i)
            hdlr(i);
    }

private:
    const int start_;
    const int end_;
};

// Coroutine と Visitor を仲介するクラス
// Coroutine は内部のループを抜けるときも戻り値を返さないといけないので
// Boost.Optional で逃げる
template<typename T_, typename Touter_>
class adapter
{
public:
    typedef boost::optional<T_> result_type;
    typedef boost::coroutines::generator<result_type> generator_type;

private:
    template<typename Tself_>
    class yielder
    {
    public:
        yielder(Tself_& gen): gen_(gen) {}

        void operator()(T_ const& i) {
            gen_.yield(i);
        }

    private:
        Tself_& gen_;
    };

public:
    template<typename Tself_>
    result_type operator()(Tself_& self) {
        yielder<Tself_> y(self);
        outer_.accept(y);
        return result_type();
    }

    adapter(Touter_& outer): outer_(outer) {}

private:
    Touter_& outer_;
};

template<std::size_t owner_arg = 1, class Base = boost::python::default_call_policies>
struct return_managed_custodian: public boost::python::with_custodian_and_ward_postcall<0, owner_arg, Base>
{
    typedef boost::python::manage_new_object result_converter;
};

// Generator コンセプトを持つクラスを Python のイテレータとして見せるための
// ラッパ
template<typename Tgen_>
class generator_iterator
{
public:
    void iter() {}

    typename Tgen_::value_type::value_type next() {
        typename Tgen_::value_type retval;
        if (!gen_ || !(retval = gen_())) {
            PyErr_SetNone(PyExc_StopIteration);
            boost::python::throw_error_already_set();
        }
        return *retval;
    }

    generator_iterator(Tgen_ const& gen): gen_(gen) {}

    static void register_class(char const* name = typeid(generator_iterator).name()) {
        using namespace boost::python;
        class_<generator_iterator>(name, no_init)
            .def("__iter__", &generator_iterator::iter,
                             return_value_policy<return_self<> >())
            .def("next", &generator_iterator::next)
            ;
    }

private:
    Tgen_ gen_;
};

// 何らかのコンテナクラス
class container
{
public:
    typedef adapter<int, internal_iterator> adapter_type;
    typedef adapter_type::generator_type generator_type;
    typedef generator_iterator<generator_type> generator_iterator_type;

    generator_iterator_type* iter() {
        return new generator_iterator_type(generator_type(adapter_type(ii_)));
    }

    container(): ii_(0, 10) {}

private:
    internal_iterator ii_;
};

BOOST_PYTHON_MODULE(coroutine_demo)
{
    using namespace boost::python;
    container::generator_iterator_type::register_class();
    class_<container>("container")
        .def("__iter__", &container::iter, return_managed_custodian<>())
        ;
}

以下実行例:

Python 2.5.4 (r254:67916, Sep 26 2009, 10:32:22) 
[GCC 4.3.4] on linux2
Type "help", "copyright", "credits" or "license" for more information.
>>> import coroutine_demo
>>> c = coroutine_demo.container()
>>> iter(c)
<coroutine_demo.18generator_iteratorIN5boost10coroutines9generatorINS0_8optionalIiEENS1_16shared_coroutineIFS4_vENS1_6detail5posix21ucontext_context_implEEEEEE object at 0x7fbe12bbe3d0>
>>> list(c)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>>> for i in c:
...     print i
... 
0
1
2
3
4
5
6
7
8
9