【Jupyter・Python】セルの実行時間に制限を設定するマジックコマンド

はじめに

競プロ用のプログラムなど制限時間を実際に設けて実行してみたいときに便利なマジックコマンドを作ってみました。

使い方

制限時間を設定したいセルの先頭に%%timelimit (制限時間[s])と付けるだけです。

%%timelimit 2

hogehoge
fugafuga
piyopiyo

マジックコマンドを登録する

マジックコマンドを使用したいノートブックで都度以下のコードを実行する
or
~\.ipython\profile_default\startup\setup.pyに以下のコードを書き込む
ことでマジックコマンドが使用可能になります。

from IPython.core.magic import register_cell_magic
from IPython.core.getipython import get_ipython

@register_cell_magic
def timelimit(t, cell):
    import threading, _thread
    assert t, 'Set the time limit.'
    def timeover_event():
        assert _thread.interrupt_main(), f'Time is up.(execution time > {t} s)'

    timer = threading.Timer(float(t), timeover_event)
    timer.start()
    get_ipython().run_cell(cell)
    timer.cancel()

仕組み

• threading.Timer()によりメインの処理とは別スレッドでタイマーを設定することで、設定した時間経過の後timeover_event()が呼び出されます。
• _thread.interrupt_main()はメインスレッドに KeyboardInterrupt を割り込ませることができます。
  → プログラムの実行は中断され、エラーメッセージが表示されます。
• get_ipython().run_cell(cell)によりセルの内容を実行します。
• 制限時間をより早くtimer.cancel()にたどり着けば、タイマーは解除されます。

実行例

%%timelimit 7

print('Start')
time.sleep(1)
print('1 s')
time.sleep(2)
print('3 s')
time.sleep(3)
print('6 s')
time.sleep(4)
print('10 s')
time.sleep(5)
print('Done!')

Start
1 s
3 s
6 s
Exception in thread Thread-24:
Traceback (most recent call last):
  File "C:\ProgramData\Anaconda3\lib\threading.py", line 932, in _bootstrap_inner
    self.run()
  File "C:\ProgramData\Anaconda3\lib\threading.py", line 1254, in run
    self.function(*self.args, **self.kwargs)
  File "C:\Users\keyta\AppData\Local\Temp/ipykernel_14348/4258713521.py", line 8, in timeover_event
AssertionError: Time is up.(execution time > 7 s)
---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_14348/601981788.py in <module>
      6 time.sleep(3)
      7 print('6 s')
----> 8 time.sleep(4)
      9 print('10 s')
     10 time.sleep(5)

KeyboardInterrupt: