Pythonにイテラブルに頼らないfor文を追加する

はじめに

この記事は東京大学電子情報工学科・電気電子工学科の後期実験「大規模ソフトウェアを手探る」のレポートとして書かれています。
今回は「py+hon」という班にて、他2人の班員とPythonをいじっていろんな機能を追加するということを行いました。
この記事ではPython3.7.9にイテラブルに頼らないfor文を追加する方法を書いています。
以下は他の班員の記事です。

環境

・OS:  Mac Catalina 10.15.7

・エディタ:  VScode 1.50.1

 

目標としたこと

Pythonのfor文はイテレータ を利用してリストなどのオブジェクトの要素を順に見ていく、というものしかなく、C++でいうところの範囲for文しか実装されていませんでした。他の言語にはあるのにpythonにはないのは気持ち悪い…ということでCの普通のfor文と同じ挙動をするfor文を追加しようと思いました。
つまり、 以下のような書き方ができるように改造するということにしました。

for i = 0, i < 10,  += 1:

Pythonの仕組み

新しい文法を追加するには内部で起こっていることを知る必要があります。
いろんなサイトを見た結果、だいたいこのようになっていると理解しました。

f:id:nawawan:20201030185315p:plain
こんな感じ

今回主に手を加えるべきファイルは図中に示した以下の4つのファイルです。

  • ast.c
  • compile.c

これらに順に実装していきます。

具体的な実装

Python/Grammar ~文法を定義~

拡張子が.cでも.pyでもないファイル。
このファイルでは文法を定義していて、ここにない記法をするとSyntaxErrorが表示されるようになっています。
ここにイテラブルに頼らないfor文としてnew_forを定義します。

Python/Grammarの72行目から74行目を見ると

if_stmt: 'if' test ':' suite ('elif' test ':' suite)* ['else' ':' suite]
while_stmt: 'while' test ':' suite ['else' ':' suite]
for_stmt: 'for' exprlist 'in' testlist ':' suite ['else' ':' suite]

というように、if文とwhile文とfor文の文法を定義しているのでこれを真似して

newfor_stmt: 'new_for' exprlist '=' test ',' test ',' augassign test ':' suite 

と実装。これは

new_for i = 0, i < 10, += 1: 

の形をとっており、「i」が「exprlist」に、一つめの「test」が「0」に、二つめの「test」が「i < 10」に、「augassign」が「+=」に、三つめの「test」が「1」に、「suite」が「for文の中身」に対応しています。

if文とかの定義を見ると「test」は比較文じゃないのか?と思うかもしれませんが、ただの数値でもいけました。このファイルの文法はかなり自由度が高いようで、今回のnew_forの定義の仕方も一例に過ぎません。

さらに、今回追加するnew_for文は複合文なので、複合文であることを示さなければなりません。具体的にはGrammarの70行目にcompound_stmtという複合文であることを示す部分があるので追加するだけです。

compound_stmt: if_stmt | while_stmt | for_stmt | newfor_stmt | try_stmt | with_stmt | funcdef | classdef | decorated | async_stmt

・Parser/Python.asdl ~関数を定義~

拡張子が.cでも.pyでもないファイルその2。

このファイルでは、テキストファイルから読み取ってできた構文木から抽象構文木に渡すときに使用される関数を定義します。34行目から38行目にかけてif文用やfor文用の関数が定義されているので、これを参考にします。

-- use 'orelse' because else is a keyword in target languages
| For(expr target, expr iter, stmt* body, stmt* orelse)
| AsyncFor(expr target, expr iter, stmt* body, stmt* orelse)
| While(expr test, stmt* body, stmt* orelse)
| If(expr test, stmt* body, stmt* orelse)

NewForを以下のように実装。

| NewFor(expr target, expr value1, expr test, operator op, expr value2, stmt* body)

これは以下のように対応しています。

new_for target = value1, test(比較), op value2:
       body(for文の中身)

ちなみに今回はnew_for文にelse機能はつけないつもりだったので、Python.asdlのwhile文やfor文に見られるorelseを実装していません。

・make regen-all ~おまじない~

PEP 306 -- How to Change Python's Grammar | Python.orgによると、Python-ast.hやらgraminit.hやらをGrammar、Python.asdlに対応するように変更しないといけないらしいのですが、正直どのファイルをどのように変更すればいいかこのサイトからはよくわかりません。しかし、ターミナルでこれを入力すると一発で自動的に必要なファイルを変更し、必要な関数やマクロなどを定義してくれます。

$ make regen-all

これでnew_forの骨格ができました。ようやくnew_for文を使ってもSyntaxErrorを吐かなくなります(中身を実装していないので、セグフォします)。

・中身の実装の前準備

make regen-allをした後にmakeをすると、Python/symtable.cとPython/compile.cのswitch文で"NewFor_kind"の場合がないと叱られてしまいます。

f:id:nawawan:20201031013124p:plain
叱られ

ということでcompile.cのswitch文にはこのように実装。

    case For_kind:
        return compiler_for(c, s);
    case NewFor_kind:
        return compiler_new_for(c, s);
    case While_kind:
        return compiler_while(c, s);

symtable.cには周りのコードを真似してこのように実装。

    case For_kind:
        VISIT(st, expr, s->v.For.target);
        VISIT(st, expr, s->v.For.iter);
        VISIT_SEQ(st, stmt, s->v.For.body);
        if (s->v.For.orelse)
            VISIT_SEQ(st, stmt, s->v.For.orelse);
        break;
    case NewFor_kind:
        VISIT(st, expr, s->v.NewFor.target);
        VISIT(st, expr, s->v.NewFor.value1);
        VISIT(st, expr, s->v.NewFor.test);
        VISIT(st, expr, s->v.NewFor.value2);
        VISIT_SEQ(st, stmt, s->v.NewFor.body);
        break;
    case While_kind:
        VISIT(st, expr, s->v.While.test);
        VISIT_SEQ(st, stmt, s->v.While.body);
        if (s->v.While.orelse)
            VISIT_SEQ(st, stmt, s->v.While.orelse);
        break;

for文の中身を表すbodyにのみVISIT_SEQ関数を使えば良いっぽい感じです。Python.asdlで定義した関数と関連があるはずですが、s->v.NewFor.opについてVISIT関数を用いなくてもうまくいきました(よくわかってません)。

※ compile.cに出てくるcompiler_new_forはこれから中身を実装する関数です。

・ast.c ~抽象構文木への変換~

ast.cは5000行もある大きなファイルですが、実装すべき内容はnew_forの構文木を受け取って、そこから意味のないものを除外して残ったものを抽象構文木として関数に渡す、というものです。

通常のfor文で説明すると、

f:id:nawawan:20201031015801p:plain
for文の構文木

このような構文木としてfor文は渡されるので、ここから意味のある「i」「range(N)」「suite」のみを抜き取って、関数に渡すということをast.cでは行っています。

while文用の関数やfor文用の関数を参考にnew_for文用の構文木から抽象構文木を構成する関数としてast_for_newfor_stmt関数を以下のように実装しました。

#if 1
static stmt_ty
ast_for_newfor_stmt(struct compiling *c, const node *n0){
    const node * const n = n0;
    asdl_seq *_target, *suite_seq;
    static operator_ty op;
    expr_ty expression1, expression2, expression3;
    expr_ty target, first;
    const node *node_target;
    /*newfor_stmt: 'new_for' exprlist '=' testlist ',' test ',' augassign  test ':' suite*/
    REQ(n, newfor_stmt);
    node_target = CHILD(n, 1);
    _target = ast_for_exprlist(c, node_target, Store);
    if(!_target)
        return NULL;
    first = (expr_ty)asdl_seq_GET(_target, 0);
    if (NCH(node_target) == 1)
        target = first;
    else
        target = Tuple(_target, Store, first->lineno, first->col_offset, c->c_arena);
    expression1 = ast_for_testlist(c, CHILD(n, 3));
    if (!expression1)
        return NULL;
    expression2 = ast_for_expr(c, CHILD(n, 5));
    if (!expression2)
        return NULL;
    op = ast_for_augassign(c, CHILD(n, 7));
    if(!op)
        return NULL;
    expression3 = ast_for_testlist(c, CHILD(n, 8));
    if(!expression3)
        return NULL;
    suite_seq = ast_for_suite(c, CHILD(n, 10));
    if (!suite_seq)
        return NULL;
    return NewFor(target, expression1, expression2, op, expression3, suite_seq, LINENO(n), n->n_col_offset, c->c_arena);
}
#endif

CHILD関数で構文木中の意味のあるノードにアクセスし、最後にPython.asdlから作られた関数に渡しています。この実装は一見大変そうに見えますが、for文用の関数などを見て何が起こっているかがおおよそわかれば同じように実装するだけなので案外楽にいけました。

・compile.c ~バイトコードの構築~

compile.c内のswitch文で用いるcompiler_new_for関数の中身を実装します。
compile.cで行っていることは、抽象構文木から得られる情報をもとに、バイトコードを構成するというものです。ここでバイトコードを構築したら、あとは他のファイルにある関数がバイトコード通りに出力をするだけなのでここさえ実装できればnew_for文の完成です。

実は今回実装するfor文は以下のwhile文とほぼ同じような挙動をするので、このwhile文と同じようなバイトコードが構築できれば勝ちです。

i = 0
while i < 10:
    i += 1

Pythonには

python3 -m dis hoge.py

とすることでhoge.pyのバイトコードを見ることができるというオプションがあるのでこれを利用します。上記のwhile文のバイトコードを見てみると以下のようになります。

f:id:nawawan:20201031032643p:plain
while文のバイトコード

このバイトコードはcompiler_while関数で生成されているので、この関数を参考にすればいいとわかります。

compiler_while関数の中身は

static int
compiler_while(struct compiler *c, stmt_ty s)
{
    basicblock *loop, *orelse, *end, *anchor = NULL;
    int constant = expr_constant(s->v.While.test);
    …
    (中略)
    …
    if (constant == -1) {
        if (!compiler_jump_if(c, s->v.While.test, anchor, 0))
            return 0;
    }
    VISIT_SEQ(c, stmt, s->v.While.body);
 …
}

となっており、変数constantが比較文の結果を保持していて、その値によってループを抜けるかどうかを判定しているので、比較が行われる前にiへの代入操作を行えばいいと考えられます。

また、VISIT_SEQ関数でbodyにアクセスすることでwhile文の中身を処理しているので、new_for文ではbodyにアクセスした後にiにopで表される演算をすればいいと考えられます。

これらを考慮して、以下のようにcompiler_new_for関数を実装。

#if 1
static int
compiler_new_for(struct compiler *c, stmt_ty s){
    basicblock *loop, *end, *anchor = NULL;
    VISIT(c, expr, s->v.NewFor.value1);
    VISIT(c, expr, s->v.NewFor.target);
    int constant = expr_constant(s->v.NewFor.test);
    if (constant == 0) {
        return 1;
    }
    loop = compiler_new_block(c);
    end = compiler_new_block(c);
    if (constant == -1) {
        anchor = compiler_new_block(c);
        if (anchor == NULL)
            return 0;
    }
    if (loop == NULL || end == NULL)
        return 0;

    ADDOP_JREL(c, SETUP_LOOP, end);
    compiler_use_next_block(c, loop);
    if (!compiler_push_fblock(c, LOOP, loop))
        return 0;
    if (constant == -1) {
        if (!compiler_jump_if(c, s->v.NewFor.test, anchor, 0))
            return 0;
    }
    VISIT_SEQ(c, stmt, s->v.NewFor.body);
    ADDOP_NAME(c, LOAD_NAME, s->v.NewFor.target->v.NameConstant.value, names);
    VISIT(c, expr, s->v.NewFor.value2);
    switch(s->v.NewFor.op){
        case 1:
            ADDOP(c, INPLACE_ADD);
            break;
        case 2:
            ADDOP(c, INPLACE_SUBTRACT);
            break;
        case 3:
            ADDOP(c, INPLACE_MULTIPLY);
            break;
        case 13:
            ADDOP(c, INPLACE_FLOOR_DIVIDE);
            break;
        default:
            PyErr_Format(PyExc_SystemError,
             "bad memberdescr type for %d", s->v.NewFor.op);
            break;
    }
    VISIT(c, expr, s->v.NewFor.target);
    ADDOP_JABS(c, JUMP_ABSOLUTE, loop);

    /* XXX should the two POP instructions be in a separate block
       if there is no else clause ?
    */
    if (constant == -1)
        compiler_use_next_block(c, anchor);
    ADDOP(c, POP_BLOCK);
    compiler_pop_fblock(c, LOOP, loop);
    compiler_use_next_block(c, end);

    return 1;
}
#endif

compiler_while関数の内容をコピペして、必要なものを追加し、不必要なものを削除したものとなっています。compiler_while関数との決定的な違いはswitch文であり、opを+=と-=と*=と//=の場合にのみ対応させてそれぞれに対応するバイトコードを構築することで、iの値の変化を加減乗除全てで可能にしています。

ここまで実装してビルドし、new_for文を使ったときにセグフォしなかったら成功です。

実際の挙動

opが加算、減算、乗算、除算の場合の4通りについて試した画像を以下に示します(読みづらいコードですみません)。
f:id:nawawan:20201031040021p:plain
個人的にi * i < 10000という条件式の部分がイテラブルに頼らないfor文ならではだなぁと思えて満足感があります。

感想

Pythonという多くの人が使っているプログラミング言語に新しい文法を付け加えたというのは自慢になるなぁと思いました。
また、膨大なソースコードの全てを完全に理解せずとも重要なポイントだけ抑えておけば手を加えることができるということを実感できたので、今後何かしらのソフトを開発する場面があったときにこの経験は生かせるだろうと思います。

今回書いたこのブログでは、一見するとエスパーして次々と一発でエラーを出すことなく実装したかのような印象を得るかもしれませんが、実際は少しいじってその時の挙動を調べるということを繰り返していました(特にPython/GrammarとPython.asdlでの実装はそうでした)。繰り返し実験していくうちになんとなく理解が深まり、エラーが消えていくという過程はまさに「手探る」という感じでとても楽しく、達成感がありました。

参考

実験のHPです。

GrammarやPython.asdlに実装する際に役立ちました。

Pythonの仕組みや改造の仕方についてとても参考になります。

Pythonの文法を変更する手順が書かれています。