Python/C API で Awaitable を作る

はじめに

今年も年末が近づいてきたので無駄に Python C/API を使ってみます。目的と手段の逆転。今回は Awaitable を、次の coroutine function spam を実行して得られる Coroutine のようなものをめざします。

import asyncio
import sys


async def spam():
    print('do something')
    ret = await asyncio.sleep(1, 'RETURN VALUE')
    return ret.lower()


async def main():
     ret = await spam()
     print(ret)


if __name__ == '__main__':
    if sys.version_info < (3, 7):
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main())
        loop.close()
    else:
        asyncio.run(main())

なにか print して、 await sleep して結果を得て、それの lower メソッドを呼んで結果を返します。 3 動作。 Awaitable はもっていないが Coroutine や Generator ならばもっている send, throw, close メソッドは今回は再現しません。

await と yield from

async def, await と書かれたコードを模倣する前に。そもそもこれはなんなのかを復習します。これらの言語機能が Python に追加提案されたときの文章が PEP 492 です。
これを読んで await とは実質 yield from だったな、と雑に思い出しました。これが用意された経緯は次のようなものだったのではないでしょうか。「Generator の仕様はもともと Coroutine として使われることを視野に入れて設計されていました。一時停止・再開ができるという特性は Iterator そのものでしたし、値を送り込むことができるという特性も PEP 342 にて追加しました。そして Python 3.4 以前まで順調に実績を積んできました。しかし Generator や Iterator と Coroutine の区別がつかないことによる問題が報告され始めたので区別がつくように Python 3.5 の言語仕様に手を入れます。 __iter__ ではなく __await__ という別のメソッドを新設します。 Generator Function 用である yield from 文も流用はせず、 Coroutine Function の中で使うための別の文 await 文を新設して区別します。」 、と。 __await__ と __iter__ という別の名前が用意されてはいますが動作としては異なるものではないのでした。 CPython の PyTypeObject 構造体上でも tp_iter と tp_as_sync.am_await が用意され別のメンバ扱いとなっています。

__await__

Awaitable に対する理解が進んだので。さっそく coroutine function spam を class 文で書き下します。 Iteratable との違いは __iter__ の代わりに __await__ を持っているところですね。とはいえ名前が違うだけで、ここから Iterator を返すという点はかわりません。

class Spam:
    def __await__(self):
        ...  # TODO: どのような Itarator を返せば spam コルーチンを模倣できる？

__iter__ と __next__

分解をつづけます。 Iterator とは __iter__ で自身を返し、 __next__ で処理を再開し次の値を作り処理を中断しつつ値を返すものです。

class _Spam:
    def __iter__(self):
        return self

    def __next__(self):
        ...  # TODO: どうすれば spam コルーチンを模倣できる？

class Spam:
    def __await__(self):
        return _Spam()

さて、「coroutine function spam を実行して得られる Coroutine と似た Awaitable」を実装するためには何が必要でしょうか。「なにか print して、 await sleep して結果を得て、それの lower メソッドを呼んで結果を返す」の 3 動作のうちどこまで処理したかの状態の保持ですね。オブジェクトに _state 属性を持たせるようにします。今回はとりあえず 0, 1, 2 の int で表現しましたが、きれいに書くのであれば Enum をもちいるのがよいでしょう。
そして状態別に処理を行う __next__ を実装します。ここで問題となるのが await もしくは yield from 。 yield from とは別の Iterator に処理を委譲するものです。別の Iterator が停止するまで繰り返し処理をしつづけて得られた値をそのまま返し続けます。これと同等の処理を実装するために動作中の別の Iterator を保持する必要があります。このため _it 属性をもたせました。この Iterator が停止したら StopIteration 例外が送られてくるので value 属性をみます。これが coroutine function や generator function の return 文に、 await, yield from 式 の値に対応しています。 lower メソッドの位置に注目を。
逆に値を返すときには StopIteration 例外に値を持たせます。
一度停止した Iterator は StopIteration 例外を返し続けないといけないという制約がありますのでその対策をします。 __next__ の末尾は raise StopIteration にします。状態を保持する属性は外部から書き換えて欲しくないものであることを示すため名前を _ で始めるようにします。この _state の初期化は、 __init__ の複数回の呼び出しへの耐性をもたせるため __new__ でおこないます。

class _Spam:
    def __new__(cls):
        obj = super().__new__(cls)
        obj._state = 0
        obj._it = None
        return obj

    def __iter__(self):
        return self

    def __next__(self):
        if self._state == 0:
            print('do something')
            self._it = asyncio.sleep(1, 'RETURN VALUE').__await__()
            self._state = 1
        if self._state == 1:
            try:
                v = next(self._it)
            except StopIteration as e:
                ret = e.value
                self._it = None
                self._state = 2
                raise StopIteration(ret.lower())
            else:
                return v
        raise StopIteration


class Spam:
    def __await__(self):
        return _Spam()

Python C/API でクラスを書こう

分解が終わったので。これを Python C/API でかきます。ここから C 言語。クラスをつくるには PyTypeObject 構造体を直接書くか PyType_Spec を定義して PyType_FromSpec を呼びます。今回は PyType_FromSpec を使う方向で。
まずは _Spam よりは簡単な Spam から。 __await__ をつくるには tp_as_sync.am_await に関数を登録します。このメソッドからは _Spam のインスタンスを返す必要があるので Spam クラスの属性に _Spam クラスを持たせるようにすることにします。クラス属性追加処理は Py_mod_exec として登録するモジュールの初期化処理にて PyType_FromSpec でのクラス作成後に行います。

typedef struct {
    PyObject_HEAD
} SpamObject;


static PyObject *
advent2019_Spam_await(SpamObject *self)
{
    PyObject *_Spam_Type = PyObject_GetAttrString((PyObject *)Py_TYPE(self), "_Spam");
    if (_Spam_Type == NULL) { return NULL; }
    PyObject *it = PyObject_CallFunction(_Spam_Type, "");
    Py_DECREF(_Spam_Type);

    return it;
}


static PyType_Slot advent2019_Spam_slots[] = {
    {Py_am_await, (unaryfunc)advent2019_Spam_await},
    {0, 0},
};


static PyType_Spec advent2019_Spam_spec = {
    .name = "advent2019.Spam",
    .basicsize = sizeof(SpamObject),
    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
    .slots = advent2019_Spam_slots,
};

static int advent2019_exec(PyObject *module) {
    int ret = -1;
    PyObject *_Spam_Type = NULL;  // TODO
    PyObject *Spam_Type = NULL;

    if (!(Spam_Type = PyType_FromSpec(&advent2019_Spam_spec))) { goto cleanup; }
    // Spam._Spam = _Spam
    if (PyObject_SetAttrString(Spam_Type, "_Spam", _Spam_Type)) { goto cleanup; }

    if (PyObject_SetAttrString(module, "Spam", Spam_Type)) { goto cleanup; }
    if (PyObject_SetAttrString(module, "_Spam", _Spam_Type)) { goto cleanup; }

    ret = 0;
cleanup:
    Py_XDECREF(_Spam_Type);
    Py_XDECREF(Spam_Type);

    if (ret) { Py_XDECREF(module); }
    return ret;
}

C/API で書く Iterator

さて、後回しにした Iterator _Spam の実装ですね。まずは _SpamObject 構造体と __new__ から。状態を保持するための何らかの値 state と asyncio.sleep().__await__ イテレータを保持する it を用意。他の PyObject* を保持することになるのでガベージコレクション機構に対応する PyObject_GC_New でメモリを確保するようにします。また、初期化が終わったら PyObject_GC_Track を呼び、ガベージコレクタに自身を登録します。

typedef struct {
    PyObject_HEAD
    unsigned char state;
    PyObject *it;
} _SpamObject;


static PyObject *
advent2019__Spam_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
{
    static char *kwlist[] = {NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "", kwlist)) {
        return NULL;
    }

    _SpamObject *obj = PyObject_GC_New(_SpamObject, type);
    if (!obj) { return NULL; }
    obj->state = 0;
    obj->it = NULL;

    PyObject_GC_Track(obj);

    return (PyObject *)obj;
}

ガベージコレクションが動いた時用の関数が必要になります。抱えているオブジェクトを手繰れるようにする traverse と、参照を破棄する clear 、自身を破棄する dealloc の 3 つです。 Py_VISIT マクロ は traverse を書くのに便利。 PyObject_GC_Track と対になる PyObject_GC_UnTrack は dealloc の開始時に呼び出すようにします。

static int
advent2019__Spam_traverse(_SpamObject *self, visitproc visit, void *arg)
{
    Py_VISIT(self->it);
    return 0;
}

static int
advent2019__Spam_clear(_SpamObject *self)
{
    Py_CLEAR(self->it);
    return 0;
}

static void
advent2019__Spam_dealloc(_SpamObject *self)
{
    PyObject_GC_UnTrack(self);
    advent2019__Spam_clear(self);
    PyObject_GC_Del(self);
}

__iter__ は自身を返すだけなので簡単ですね。参照カウントの操作を忘れないようにしつつ。

static PyObject *
advent2019__Spam_iter(_SpamObject *self)
{
    Py_INCREF(self);
    return (PyObject *)self;
}

最後に __next__ 。 Python C/API 上では iternext という名前になります。組み込み関数 print と asyncio モジュールの sleep はクラス作成時に _print と _sleep という名前でクラス属性に持たせるようにしておきます、 Spam._Spam と同じ要領で。
あとは Python で書いたコードを PyObject_GetAttrString, PyObject_CallFunction らを用いて地道に移植していく作業です。呼ぶたびに戻り値が NULL になっていないかを確認します。不必要となったオブジェクトの参照カウントを減らす処理がしんどい。
try 文の移植には PyErr_Fetch, PyErr_GivenExceptionMatches, PyErr_Restore を使います。

static PyObject *
advent2019__Spam_iternext(_SpamObject *self)
{
    if (self->state == 0) {
        // print('do something')
        PyObject *printfunc = PyObject_GetAttrString((PyObject *)Py_TYPE(self), "_print");
        if (!printfunc) { return NULL; }
        PyObject *ret = PyObject_CallFunction(printfunc, "s", "do something");
        Py_DECREF(printfunc);
        if (!ret) { return NULL; }
        Py_DECREF(ret);

        // self._it = asyncio.sleep(1, 'RETURN VALUE').__await__()
        PyObject *sleep_cofunc = PyObject_GetAttrString((PyObject *)Py_TYPE(self), "_sleep");
        if (!sleep_cofunc) { return NULL; }
        PyObject *sleep_co = PyObject_CallFunction(sleep_cofunc, "is", 1, "RETURN VALUE");
        Py_DECREF(sleep_cofunc);
        if (!sleep_co) { return NULL; }
        if (!(Py_TYPE(sleep_co)->tp_as_async)) { Py_DECREF(sleep_co);  return NULL; }
        if (!(Py_TYPE(sleep_co)->tp_as_async->am_await)) { Py_DECREF(sleep_co);  return NULL; }
        PyObject *temp = self->it;
        self->it = Py_TYPE(sleep_co)->tp_as_async->am_await(sleep_co);
        Py_DECREF(sleep_co);
        Py_XDECREF(temp);
        if (self->it == NULL) { return NULL; }

        self->state = 1;
    }
    if (self->state == 1) {
        // next(self.it)
        if (Py_TYPE(self->it)->tp_iternext == NULL) { PyErr_SetString(PyExc_TypeError, "no iternext"); return NULL; }
        PyObject *ret = Py_TYPE(self->it)->tp_iternext(self->it);
        if (!ret) {
            // except StopIteration as e
            PyObject *type, *value, *traceback;
            PyErr_Fetch(&type, &value, &traceback);
            if (PyErr_GivenExceptionMatches(type, PyExc_StopIteration)) {
                Py_XDECREF(type);
                Py_XDECREF(traceback);
                if (!value) { PyErr_SetString(PyExc_ValueError, "no StopIteration value"); return NULL; }
                // ret = e.value.lower()
                PyObject *value2 = PyObject_CallMethod(value, "lower", NULL);
                Py_DECREF(value);
                if (!value2) { return NULL; }
                // raise StopIteration(ret)
                PyErr_SetObject(PyExc_StopIteration, value2);
                Py_DECREF(value2);

                Py_CLEAR(self->it);
                self->state = 2;
            } else {
                // except:
                //     raise
                PyErr_Restore(type, value, traceback);
            }
        }
        return ret;
    }

    // raise StopIteration(None)
    PyErr_SetNone(PyExc_StopIteration);
    return NULL;
}

これで _Spam クラスのメソッドがそろったので PyType_Spec を定義します。ガベージコレクションで管理されるべきクラスであることををしめすフラグ Py_TPFLAGS_HAVE_GC を設定するようにします。

static PyType_Slot advent2019__Spam_slots[] = {
    {Py_tp_new, advent2019__Spam_new},
    {Py_tp_iter, advent2019__Spam_iter},
    {Py_tp_iternext, advent2019__Spam_iternext},
    {Py_tp_traverse, advent2019__Spam_traverse},
    {Py_tp_clear, advent2019__Spam_clear},
    {Py_tp_dealloc, advent2019__Spam_dealloc},
    {0, 0},
};


static PyType_Spec advent2019__Spam_spec = {
    .name = "advent2019._Spam",
    .basicsize = sizeof(_SpamObject),
    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
    .slots = advent2019__Spam_slots,
};

おわりに

Python C/API で Awaitable を実装してみました。また、これをおこなうにあたって必要な情報がまとまっている公式ドキュメントへのリンクを集めることができました。
ところで。これは役に立つのでしょうか？ たったこれだけのことのためのコードがずいぶんと長く、 async def, await 構文の便利さを思い知っただけのような、いや CPython への理解を深めるのには役立つかもしれない………。

setup.cfg

[metadata]
name = advent2019
version = 0.0.0

[options]
python_requires = >=3.5.0

setup.py

from setuptools import Extension, setup

extensions = [Extension('advent2019', sources=['advent2019.c'])]

setup(ext_modules=extensions)

advent2019.c

#define PY_SSIZE_T_CLEAN
#include <Python.h>


typedef struct {
    PyObject_HEAD
    unsigned char state;
    PyObject *it;
} _SpamObject;


static PyObject *
advent2019__Spam_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
{
    static char *kwlist[] = {NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "", kwlist)) {
        return NULL;
    }

    _SpamObject *obj = PyObject_GC_New(_SpamObject, type);
    if (!obj) { return NULL; }
    obj->state = 0;
    obj->it = NULL;

    PyObject_GC_Track(obj);

    return (PyObject *)obj;
}


static PyObject *
advent2019__Spam_iter(_SpamObject *self)
{
    Py_INCREF(self);
    return (PyObject *)self;
}


static PyObject *
advent2019__Spam_iternext(_SpamObject *self)
{
    if (self->state == 0) {
        // print('do something')
        PyObject *printfunc = PyObject_GetAttrString((PyObject *)Py_TYPE(self), "_print");
        if (!printfunc) { return NULL; }
        PyObject *ret = PyObject_CallFunction(printfunc, "s", "do something");
        Py_DECREF(printfunc);
        if (!ret) { return NULL; }
        Py_DECREF(ret);

        // self._it = asyncio.sleep(1, 'RETURN VALUE').__await__()
        PyObject *sleep_cofunc = PyObject_GetAttrString((PyObject *)Py_TYPE(self), "_sleep");
        if (!sleep_cofunc) { return NULL; }
        PyObject *sleep_co = PyObject_CallFunction(sleep_cofunc, "is", 1, "RETURN VALUE");
        Py_DECREF(sleep_cofunc);
        if (!sleep_co) { return NULL; }
        if (!(Py_TYPE(sleep_co)->tp_as_async)) { Py_DECREF(sleep_co);  return NULL; }
        if (!(Py_TYPE(sleep_co)->tp_as_async->am_await)) { Py_DECREF(sleep_co);  return NULL; }
        PyObject *temp = self->it;
        self->it = Py_TYPE(sleep_co)->tp_as_async->am_await(sleep_co);
        Py_DECREF(sleep_co);
        Py_XDECREF(temp);
        if (self->it == NULL) { return NULL; }

        self->state = 1;
    }
    if (self->state == 1) {
        // next(self.it)
        if (Py_TYPE(self->it)->tp_iternext == NULL) { PyErr_SetString(PyExc_TypeError, "no iternext"); return NULL; }
        PyObject *ret = Py_TYPE(self->it)->tp_iternext(self->it);
        if (!ret) {
            // except StopIteration as e
            PyObject *type, *value, *traceback;
            PyErr_Fetch(&type, &value, &traceback);
            if (PyErr_GivenExceptionMatches(type, PyExc_StopIteration)) {
                Py_XDECREF(type);
                Py_XDECREF(traceback);
                if (!value) { PyErr_SetString(PyExc_ValueError, "no StopIteration value"); return NULL; }
                // ret = e.value.lower()
                PyObject *value2 = PyObject_CallMethod(value, "lower", NULL);
                Py_DECREF(value);
                if (!value2) { return NULL; }
                // raise StopIteration(ret)
                PyErr_SetObject(PyExc_StopIteration, value2);
                Py_DECREF(value2);

                Py_CLEAR(self->it);
                self->state = 2;
            } else {
                // except:
                //     raise
                PyErr_Restore(type, value, traceback);
            }
        }
        return ret;
    }

    // raise StopIteration(None)
    PyErr_SetNone(PyExc_StopIteration);
    return NULL;
}


static int
advent2019__Spam_traverse(_SpamObject *self, visitproc visit, void *arg)
{
    Py_VISIT(self->it);
    return 0;
}

static int
advent2019__Spam_clear(_SpamObject *self)
{
    Py_CLEAR(self->it);
    return 0;
}

static void
advent2019__Spam_dealloc(_SpamObject *self)
{
    PyObject_GC_UnTrack(self);
    advent2019__Spam_clear(self);
    PyObject_GC_Del(self);
}


static PyType_Slot advent2019__Spam_slots[] = {
    {Py_tp_new, advent2019__Spam_new},
    {Py_tp_iter, advent2019__Spam_iter},
    {Py_tp_iternext, advent2019__Spam_iternext},
    {Py_tp_traverse, advent2019__Spam_traverse},
    {Py_tp_clear, advent2019__Spam_clear},
    {Py_tp_dealloc, advent2019__Spam_dealloc},
    {0, 0},
};


static PyType_Spec advent2019__Spam_spec = {
    .name = "advent2019._Spam",
    .basicsize = sizeof(_SpamObject),
    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
    .slots = advent2019__Spam_slots,
};


typedef struct {
    PyObject_HEAD
} SpamObject;


static PyObject *
advent2019_Spam_await(SpamObject *self)
{
    PyObject *_Spam_Type = PyObject_GetAttrString((PyObject *)Py_TYPE(self), "_Spam");
    if (_Spam_Type == NULL) { return NULL; }
    PyObject *it = PyObject_CallFunction(_Spam_Type, "");
    Py_DECREF(_Spam_Type);

    return it;
}


static PyType_Slot advent2019_Spam_slots[] = {
    {Py_am_await, (unaryfunc)advent2019_Spam_await},
    {0, 0},
};


static PyType_Spec advent2019_Spam_spec = {
    .name = "advent2019.Spam",
    .basicsize = sizeof(SpamObject),
    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
    .slots = advent2019_Spam_slots,
};


static int advent2019_exec(PyObject *module) {
    int ret = -1;
    PyObject *builtins = NULL;
    PyObject *printfunc = NULL;
    PyObject *asyncio_module = NULL;
    PyObject *sleep = NULL;
    PyObject *_Spam_Type = NULL;
    PyObject *Spam_Type = NULL;

    if (!(builtins = PyEval_GetBuiltins())) { goto cleanup; }  /* borrowed */
    // fetch the builtin function print
    if (!(printfunc = PyMapping_GetItemString(builtins, "print"))) { goto cleanup; }

    // import asyncio
    if (!(asyncio_module = PyImport_ImportModule("asyncio"))) { goto cleanup; }
    if (!(sleep = PyObject_GetAttrString(asyncio_module, "sleep"))) { goto cleanup; };

    if (!(_Spam_Type = PyType_FromSpec(&advent2019__Spam_spec))) { goto cleanup; }
    // _Spam._print = print
    if (PyObject_SetAttrString(_Spam_Type, "_print", printfunc)) { goto cleanup; }
    // _Spam._sleep = asyncio.sleep
    if (PyObject_SetAttrString(_Spam_Type, "_sleep", sleep)) { goto cleanup; }

    if (!(Spam_Type = PyType_FromSpec(&advent2019_Spam_spec))) { goto cleanup; }
    // Spam._Spam = _Spam
    if (PyObject_SetAttrString(Spam_Type, "_Spam", _Spam_Type)) { goto cleanup; }

    if (PyObject_SetAttrString(module, "Spam", Spam_Type)) { goto cleanup; }
    if (PyObject_SetAttrString(module, "_Spam", _Spam_Type)) { goto cleanup; }

    ret = 0;
cleanup:
    Py_XDECREF(printfunc);
    Py_XDECREF(asyncio_module);
    Py_XDECREF(sleep);
    Py_XDECREF(_Spam_Type);
    Py_XDECREF(Spam_Type);

    if (ret) { Py_XDECREF(module); }
    return ret;
}


static PyModuleDef_Slot advent2019_slots[] = {
    {Py_mod_exec, advent2019_exec},
    {0, NULL}
};


static struct PyModuleDef advent2019_moduledef = {
    PyModuleDef_HEAD_INIT,
    .m_name = "advent2019",
    .m_slots = advent2019_slots,
};


PyMODINIT_FUNC PyInit_advent2019(void) {
    return PyModuleDef_Init(&advent2019_moduledef);
}

これをつかってみるコードの例

import sys
import asyncio

import advent2019


async def main():
    v = await advent2019.Spam()
    print(v)


if __name__ == '__main__':
    if sys.version_info < (3, 7):
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main())
        loop.close()
    else:
        asyncio.run(main())