複雑な内包表記やreduceを読みやすく書くためのPythonライブラリを作った

複雑な内包表記を読みやすく書くためのTipsみたいな記事を前に書いたが、それを元にxcompというライブラリを作った。Githubで公開中。

以下が基本的な目標。

• 通常のfor文の語順で、
• インデントも入れつつ、
• 状態変化なしで
• コレクションを初期化する

インストール

pip install git+git://github.com/ukyo-su/xcomp@master

comp(collector)

for xcomp import comp

# a = [i for i in range(5)]の代わりに
@comp(list)
def a():
    for i in range(5):
        yield i

# a == [0, 1, 2, 3, 4]

# 和も取れる。dictとかsetも可能
@comp(sum)
def a():
    for i in range(5):
        yield i

# a == 10

# 複雑な内包も読みやすく書ける
@comp(list)
def fizz_buzz():
    for i in range(30):
        if i % 15 == 0:
            yield "Fizz Buzz"
        elif i % 3 == 0:
            yield "Fizz"
        elif i % 5 == 0:
            yield "Buzz"
        else:
            yield i

multi_comp(*collectors)

一回のループで複数のコレクションを作りたい場合。

from xcomp import multi_comp

data = [(0, 1), (2, 3)]

@multi_comp(list, list)
def a():
    for i, j in data:
        yield i, j

a_i, a_j = a

# a_i == [0, 2]
# a_j == [1, 3]

reduce_comp(init, for_=None, for_nest=(), result=lambda x: x)

functools.reduceを少し読みやすく書くためのデコレータ。
Racketのfor/foldが元ネタ。

from xcomp import reduce_comp

# reduce(lambda sum_, i: sum_ + i, range(5), 0)
@reduce_comp(0, for_=range(5))
def a(sum_, i):
    return sum_ + i

# a == 10

# ネストしたforもOK
@reduce_comp([], for_nest=(range(3), range(2)))
def a(acc, i, j):
    return [*acc, (i, j)]

# a == [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]

# 最後に一手間加える場合
@reduce_comp([], for_=range(5),
             result=lambda x: list(reversed(x)))
def a(acc, i):
    return [*acc, i]

# a == [4, 3, 2, 1, 0]

# break, continueも対応
@reduce_comp(0, for_=range(100))
def a(acc, i):
    if i > 5:
        raise BreakReduce
    return acc + 1

# a == 6

@reduce_comp(0, for_=range(10))
def a(acc, i):
    if i < 5:
        raise ContinueReduce
    return acc + i

# a == 35

（本当に読みやすいのか……？）

multi_reduce_comp(*inits, for_=None, for_nest=(), result=lambda *args: args)

累積していく変数が複数ある場合のreduce_comp。Racketのfor/foldにはこっちが近い。

@multi_reduce_comp(0, [],
                   for_=range(5))
def a(sum_, rev_list, i):
    return sum_ + i, [i, *rev_list]

a_sum, a_rev_list = a

# a_sum == 10
# a_rev_list == [4, 3, 2, 1, 0]

@multi_reduce_comp([], set(),
                   for_=[0, 1, 1, 2, 3, 4, 4, 4],
                   result=lambda acc, seen: list(reversed(acc)))
def a(acc, seen, i):
    if i in seen:
        return acc, seen
    else:
        return [i, *acc], {*seen, i}

# a == [0, 1, 2, 3, 4]

delay_arg(func, *args, **kwargs)

一つ目の引数を最後に回すための高階関数。delay_arg(f, *args, **kwargs)(a)はf(a, *args, **kwargs)になる。

map, filter, reduceなんかをデコレータとして運用可能。

@list
@delay_arg(map, range(5))
def a(i):
    return i * 2

# a == [0, 2, 4, 6, 8]