💭

簡易tqdm関数作ってみた(Python)

に公開

概要

こんにちは!本日は簡易的なtqdmを自作してみたので、考え方やコードをまとめてみました!
ターミナル上で視覚的に動きがあるとワクワクするものがありますね!

目次

  1. tqdmってなに?
  2. 自作パート(関数定義)
  3. 自作パート(内部コード)
  4. まとめ

1. tqdmってなに?

tqdmアラビア語の「taqaddum (تقدّم)」に由来し、「進歩」を意味し、スペイン語では「I love you so much ( te quiero demasiado ) 」の略語です。
tqdmの公式githubより

tqdmとはPythonで使われるプログレスバーのことです
動かしてみました

pip install tqdm

簡単なコードを書くだけ

from tqdm import tqdm
import time

for i in tqdm(range(100)):
    time.sleep(0.5)

結果

100%|████████████████████████████| 100/100 [00:30<00:00,  3.29it/s]

easyですね!sleep(x)xの値を調整することで更新速度を調整できます。

中身の部分に触れていきましょう!
まずtqdm()反復可能なオブジェクト(iterable) を受け取ります
返り値はイテレータをラップして、回すたびに表示するオブジェクト が返ります

tqdm(["a", "b", "c", "d"])

これもできるってことですね
https://github.com/tqdm/tqdm

2. 自作パート(関数定義)

結論から

def ft_tqdm(lst: range)

今回、iterableだと、

  • 長さがわからない
  • 何個来るかが途中まで不明

というケースも含まれてしまうので、仕組みを分かりやすくするため、入力は range に限定しています。
range は開始・終了・ステップが明確で、進捗率を計算しやすいためです。

3. 自作パート(内部コード)

今回のコードを一言で表すのなら、for で range を 1 要素ずつ進めながら、その進捗に応じて「進捗バーっぽい文字列」を表示する です
簡単そうに聞こえますよね
今回のポイントは2つあります.

  • generator
  • yield

です。

generator とは値を一気に返すのではなく、必要になった分だけ、1つずつ返す関数 です
普通の関数

def normal():
    return 1
  • 呼ばれた瞬間に処理が全部終わる
  • returnしたら関数は終了

generator

def gen():
    yield 1
    yield 2
    yield 3
  • 呼ばれただけでは何も起きない
  • for文で回されるたびに1つずつ値を返す
  • 途中状態を覚えたまま再開する

tqdmはgenerator の動きを作ろうという話だと理解できますね
ここでどう呼び出すかを考えましょう!

for elem in ft_tqdm(range(333)):
    sleep(0.005)

と呼び出したいですよね
そうすると関数側では、

def ft_tqdm(lst: range):
    """Recode tqdm"""
    for elem in lst:
        # 進捗表示
        yield elem

これで一旦最低限呼び出された時に何を返すか という部分はできましたね

次は進捗をターミナル上に表示できるようにしたいです
やりたことはシンプルです

forで1要素進むたびに
現在どこまで進んだかを計算して
同じ行に上書きする

これだけです。
進捗バーの表示に必要な最低限の情報は、

  • 全体の長さ
  • 今何個目か
  • 進捗率
  • バーの見た目用の長さ
  1. 全体の長さ、今何個目か
    今回は入力がrangeに限定されているので、全体の長さは簡単ですね
total = len(lst)

そしてfor文で回しているので今何個目かはカウンタで管理できます。

count += 1

これでcountが現在の進捗
totalが全体になりますね。

  1. 進捗率は
prog = count / total

パーセント表示なら

percent = prog * 100

ですね

  1. 進捗バー的なのを作る
    今回get_terminal_size()でターミナルの横幅を取得し、可変のサイズにしたいと思います。
term_width = get_terminal_size().columns

最低限の幅を確保したいのでガードを入れます

bar_len = max(10, bar_len)

バーの長さを設定したら、全体の長さの進捗率分だけ表示し、それ以外はスペースで埋めます。

filled_len = int(bar_len * prog)

表示する進捗は、

bar = "█" * filled_len + " " * (bar_len - filled_len)

もう一つポイントがあります
ターミナルで同じ行を更新するには行頭に戻る制御文字を使います。

\r

さらに、出力を即座に反映させたいのでflushも有効にする
まとめると

print(f"\r{percent:6.2f}%|{bar}| {i}/{total}", end="", flush=True)

4. まとめ

今回は進捗状況をターミナル上に出力できるtqdm関数を簡易的に自作してみました。

from shutil import get_terminal_size

def ft_tqdm(lst: range):
    total = len(lst)
    term_width = get_terminal_size().columns
    bar_len = max(10, term_width - 25)

    count = 0
    for elem in lst:
        count += 1
        progress = count / total
        filled = int(bar_len * progress)

        bar = "█" * filled + " " * (bar_len - filled)
        percent = progress * 100

        print(f"\r{percent:3d}%|{bar}| {count}/{total}", end="", flush=True)
        yield elem

    print()

for + printで連続で出力するよりも
実際のツールっぽくなりますね!

Discussion