asciitree の AttributeTraversal を使う

pypi に asciitree というパッケージがある。 ターミナルにツリー表示を描くために使用することができるのだけど、README をパッと見る限りではツリー構造を dict で用意しなければならないようで、なんだか不便そう。ただ、ドキュメントを見ると children という属性を見てツリー構造を解釈してくれそうな AttributeTraversal というクラスがあるが、使用例がない。

というわけで、その使用例を書いてみる。書き出すツリーは、README と同じ以下のもの。

asciitree
 +-- sometimes
 |   +-- you
 +-- just
 |   +-- want
 |       +-- to
 |       +-- draw
 +-- trees
 +-- in
     +-- your
         +-- terminal

まずは、import。

from asciitree import LeftAligned
from asciitree.traversal import AttributeTraversal

そして、children を属性に持つ Node クラスを用意してみる。（とりあえず tuple にしているけど、list でも set でも良い。） この Node クラスのオブジェクトに対して str() が行われた結果が、ツリーの各ノードに表示される文字列になるので、__repr()__ でカスタマイズできる。

class Node:
    def __init__(self, data):
        self.__data = data
        self.children = tuple()

    def __repr__(self):
        return str(self.__data)

ノードを作っていく。

ast = Node('asciitree')
sts = Node('sometimes')
you = Node('you')
jst = Node('just')
wnt = Node('want')
to  = Node('to')
drw = Node('draw')
trs = Node('trees')
inn = Node('in')
yr  = Node('your')
tnl = Node('terminal')

各ノードが持つ子どもを tuple にまとめて、children に入れていく。（list とか set にする場合には、それに応じて適当に。）

ast.children = (sts, jst, trs, inn)
sts.children = (you,)
jst.children = (wnt,)
wnt.children = (to, drw)
inn.children = (yr,)
yr.children  = (tnl,)

traverse 引数に AttributeTraversal() を指定して LeftAligned のオブジェクトを作り、call する。

tr = LeftAligned(traverse=AttributeTraversal())
print(tr(ast))

これで、ツリー構造をターミナルに書くことができる。 dict で表現するよりはかなり楽。

ちなみに、asciitree はデフォルトで children という名前の属性を見てくれるのだけど、これを変更することも可能。AttributeTraversal オブジェクトを作る際に、attribute 引数で属性名を文字列で指定すれば良い。例えば、子どもがひとつしかない構造で属性名が child があるというような場合には、AttributeTraversal(attribute='child') と書けば良い。