Jupyter Notebook/Labでopen(0)を使う（JupyterでAtCoderの問題を効率よくやるために）

最近AtCoder(競技プログラミング)をPythonで解いてて他者のコードを参考にしつつ、Jupyter Notebook/Labで実現できるか試してみたところちょっとハマったのと、できたものが便利だったので忘備録的に残しておこうと思います。
最初はopen(0)についての説明になります。open(0)が何か知っていて早く本題を読みたい方は飛ばして下さい。

Pythonで標準入力から受け取る

※この辺りは必要最低限の説明で、競プロ等で使われる標準入力の方法としてはこちらの記事が詳しいのでそちらをご覧になるのが良いかと思います。
※当章の内容はIPython上での動作になります。

通常Pythonで標準入力から値や文字列を受け取る時は以下のようにします。

>>> a = input()
test
>>> print(a)
'test'

これは自明ですね。
では、AtCoderでよくある複数数値、例えばこの問題のように複数の入力があった場合はどうしているかというと・・・


よくある入力例としてはこんな感じ

W, H, N = map(int, input().split())
S = [list(map(int, input().split())) for _ in range(N)]

念の為確認すると以下のとおりになります。

>>> W, H, N = map(int, input().split())
5 4 2
>>> S = [list(map(int, input().split())) for _ in range(N)]
2 1 1
3 3 4
>>> print(W, H, N)
5 4 2
>>> print(S)
[[2, 1, 1], [3, 3, 4]]

WとHとNはinputで取得した文字列を空白で切ってリスト化した後、mapでintにキャストして、それぞれWとHとNに入れています。
Sは1行目で取得したN行分だけinput()を繰り返し、同様に文字列を空白で区切ってintにキャストした値をリストに順次入れていくというものです。

open(0)を使った入力

今回は上記のでも良いのですが、もっと一度に入力したいケースが出て来ます。
またinput()を複数回繰り返すのはより短いコードで実現したい場合に不利になります。
ここでopen(0)を使った例を示します。

*D, = open(0)
もしくは
D = open(0).read()

openは組込み関数でファイルを読込むものですが、0や1を指定すると標準入力や標準出力が対象になります（参考）。
open()の引数に0を指定すると標準入力からの入力になるので、一度に複数の入力値を標準入力から得られます。

1行毎に値をリストに取得したい場合は"*"と","を忘れないで下さい。
open(0).read()とすると、改行も含め文字列としてDに格納されます。

>>> *D, = open(0)
5 4 2
2 1 1
3 3 4
^Z
>>> print(D)
['5 4 2\n', '2 1 1\n', '3 3 4\n']

もちろんこのままでは全て文字列型で値として扱いたいので以下のようにmapでintにキャストしたものを利用することが多いかと思います。

>>> D = [list(map(int, s.split())) for s in open(0)]
5 4 2
2 1 1
3 3 4
^Z
>>> print(D)
[[5, 4, 2], [2, 1, 1], [3, 3, 4]]

なお、1行目のW, H, Nを変数にして、それ以降の行をデータとして取得したい場合、さらには以下のように改良できます。

>>> (W, H, N), *D, = [list(map(int, s.split())) for s in open(0)]
5 4 2
2 1 1
3 3 4
^Z
>>> print(W, H, N)
5 4 2
>>> print(D)
[[2, 1, 1], [3, 3, 4]]

open(0)の説明としてはちょっと長くなりましたが、こんな感じで競プロの入力方法としてopen(0)は結構使われるのを見ます。

Jupyter Notebook/Labでopen(0)を使う

他の回答者の回答を一つずつ試して比べるのにJupyterLabを使っているのですが、短いコードにはよくopen(0)が使われていることがままあります。そこでJupyter NotebookやLab上でopen(0)を実行すると以下のようになります。

という感じでエラーが出て、そのままでは実行できません。

Jupyter Notebook/Labでopen(0)を実現するコード

Jupyterで実施できるようopen(0)をinput()を使った処理に置き換えて実行しても良いのですが、折角の先人のコードをそのまま実行したいのと、いちいち毎行Jupyter Notebook/Lab上にサンプル入力をコピペするのは面倒だったり、あとできればopen(0)で記載しデバッグが通った内容を提出したいと思ったので色々と調べていたところ良い記事がありました。

　　　AtCoderのためにjupyterの入力を複数行に対応させる

上記記事はopen(0)に関するものではなく複数行入力できる内容だったのですが、こちらをもう一つ味付けすれば実現できそうだと思ったので、記事を参考にしつつ、そのままopen(0)をできるようにしたのが以下の内容です。

from ipywidgets import Textarea
import io

if 'open' in globals():
    del open
original_open = open

class custom_open():
    def __init__(self):
        self.text = ''

    def __call__(self, file, *args, **kwargs):
        if file == 0:
            return io.StringIO(self.text)
        return original_open(file, *args, **kwargs)

    def updater(self, change):
        self.text = change["new"]

open = custom_open()
text_area = Textarea()
text_area.observe(open.updater, names='value')
display(text_area)

ファイル読込み機能を削除したopen関数(旧版)はここをクリックして下さい。

こちらは要は組込み関数openを「lambda式でテキストボックスの内容を取得する」というものに置き換えた内容になります。

from ipywidgets import Textarea
import io

def updater(change):
    global open
    open = lambda x: io.StringIO(change["new"])

text_area = Textarea()
text_area.observe(updater, names='value')
display(text_area)

こちらのコードの場合、無理矢理open関数をlambdaで置き換えただけで、"open関数をjupyter Notebook/Labでテキストボックスの内容を取得する関数として使う"と割り切ったものなので、ファイル入力としてのopen関数の機能は無くなっていることに注意下さい。元のファイル入出力を行うopen関数に戻すには後述の通り「del open」を実施ください。
ファイル入出力機能を残したままのopen関数として使いたい場合は上記のコードを仕様下さい。

上記コードを実行すると以下のようなテキストボックスが現れます。

その後、別のセル上でopen(0)を実施すると、表示されたテキストボックスの内容を(open(0)を使う感覚で)取得できます。
このコードは各Notebookファイルの冒頭で実行しておけばよく、テキストボックスの入力内容が更新されれば次回実行時のopen(0)で得られる内容も変化します。（再度上記コードを入力したセルを実行する必要ない）

上記のopen(0)で得られるものはStringIOオブジェクトなのでopen(0).read()やopen(0).readline()も動作します。
もちろんファイルの読込みも元のopen関数同様open('ファイルパス')で実施できます。

こちらを使えば例えばこのAtCoderの問題において、「コード長が一番短くopen(0)を使った回答」をJupyter Notebook/Lab上で動作確認することができます。

なお、open関数を元の組込み関数に戻したい場合は、以下のコードを実行すれば、以降はただのファイル読込み関数としてのopen関数が利用できるようになります。

del open

JupyterLabで実施する場合の注意

上記のコードはJupyter Notebook上では基本的に動くと思いますが、JupyterLabにおいては初期設定のままではテキストボックスがでてこず動きませんので、以下を実施し拡張機能を有効化し必要なExtensionを入れる必要があります（参考）。

1. 拡張機能を有効にする（有効にしてあれば飛ばして下さい）

   1. Node.jsをインストールする
   2. JupyterLabを開き、メニューバーより[Settings] → [Advanced Settings Editor]をクリック
   3. 左タブの「Extension Manager」をクリックし、"enabled": trueにして右上の保存(フロッピーディスク)アイコンをクリック

2. widgets用のExtensionを入れる
   拡張機能を有効にしたら以下2つのいずれかの方法でjupyter-widgetsを入れて下さい。

   • コマンドプロンプト/シェル上で入れる
     • JupyterLabを終了させ、以下のコマンドを実行

   jupyter labextension install @jupyter-widgets/jupyterlab-manager
   

   • JupyterLabで入れる
     • JupyterLabを開き、「Extension Manager」アイコン(ジグソーパズルのピース形状)をクリック
     • 「SEARCH」テキストボックスに「widgets」と入れる
     • SEARCH RESULTを開き「jupyter-widgets/jupyterlab-manager」を見つけ「install」をクリック
     • 途中「REBUILD」するかどうかを聞かれたら「REBUILD」をクリックする。
     • インストールが終わったらRELOADをクリックする
       （ここまでに結構時間かかる。インストールの進捗はシェル上で確認できる）

上記Extensionを入れて再度open(0)を有効にするコードを実行すればJupyterLabでも使えるようになっているはずです。

追記（inputもテキストボックスから取得する）

open(0)だけでなく、どうせならinputもテキストボックスから取得できるようにコードを改良しました。
こちらであればinput関数もテキストボックス上からデータを取得できるようになります。
[input() for _ in range(N)]も問題なく動作しますので、inputとopen(0)で入力欄を使い分けなくてすむかと思います。
もちろんopen(0)もJupyter上で使えます（Jupyter Notebook/Lab両方で動作確認済み）

from ipywidgets import Textarea
import io

if 'open' in globals():
    del open
if 'input' in globals():
    del input

original_open = open

class custom_open():
    def __init__(self):
        self.text = ''

    def __call__(self, file, *args, **kwargs):
        if file == 0:
            return io.StringIO(self.text)
        return original_open(file, *args, **kwargs)

    def updater(self, change):
        self.text = change["new"]

class custom_input():
    def __init__(self):
        self.__sio = io.StringIO('')
        self.shell = get_ipython()
        if self.shell.events.callbacks['pre_run_cell'] != []:
            self.shell.events.callbacks['pre_run_cell'] = []
        self.shell.events.register('pre_run_cell', self.pre_run_cell)

    def __call__(self):
        return self.__sio.readline().strip()

    def pre_run_cell(self, info):
        text = self.shell.user_ns.get('text_area', None).value
        self.__sio = io.StringIO(text)

open = custom_open()
input = custom_input()

text_area = Textarea()
text_area.observe(open.updater, names='value')
display(text_area)

こちらも組込み関数inputおよびopenを置き換えてますので、元に戻す場合は「del input」「del open」して下さい。

まとめ

Jupyter Notebook/Labでopen(0)を使えないのは有名ですが、このようになんとか無理矢理？実現することは可能です。
実際、複数行の入力をただJupyterで実現するだけであれば元記事の方法で充分かと思います。
一方で当方法のメリットとしては

• Jupyter上でAtCoderのコードテストページのように動作させられる
  （サンプル入力をコピーし、テキストボックスに入力してそのまま実行できる）
• AtCoderで他参加者がopen(0)を用いた入力方式で書いたコードをJupyter上でそのまま動かせる
• 各セルにおいてopen(0)で記載したコードをAtCoderやIPython上にそのまま持って行っても動く

と言う点が挙げられます。

競プロをJupyter Notebook/Labで実施するのはマイナーかもしれませんが、あれこれ勉強中であれば今回試した方法はそれなりに便利なので参考になればと思います。
なによりJupyter Notebook/Lab上でopen(0)の他者の素晴らしいコードをそのまま実施できたり、input()で一行ずつ入力しなくても済むようになるのは結構便利です。参考にさせて頂いた各記事の皆さんに多謝。
ともあれ、少しでもお役に立てればと思います。