🐍

Visual Studio Code の Python 拡張で日本語を含むコードが REPL で実行されないバグを修正してみた

2024/11/22に公開

Background

ある日 Visual Studio Code で python を書いていて、 Shift + Enter で実行できる REPL が特定の条件で動作しないことに気づきました。 (ずっとグルグルが出ています)

この記事ではどうやってこのバグを調査し、修正したかを記載します。

バグの再現確認と状況の確認

色々試していたところ、下記の条件でバグが顕在化することがわかりました。

  • 日本語 (non-ascii の文字、 utf-8 で 2 byte 以上で表現される文字) がコードに含まれている

MacOS・Linux:

  • Visual Studio Code 1.95.3 pre-release 版でも再現。
  • Python Extension 2024.20.0

issues もいくつかのクエリで確認しましたが、関連するバグは見つけられませんでした。

バグの調査

まず、 Visual Studio Code の開発ツールを見てみたところ、以下のようなエラーが出ていました

[Extension Host] Log: Traceback (most recent call last):
  File "/home/tomoki/.vscode/extensions/ms-python.python-2024.20.0-linux-x64/python_files/python_server.py", line 183, in <module>
    request_json = json.loads(request_text)
                   ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/json/__init__.py", line 346, in loads
    return _default_decoder.decode(s)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/json/decoder.py", line 341, in decode
    raise JSONDecodeError("Extra data", s, end)
json.decoder.JSONDecodeError: Extra data: line 1 column 111 (char 110)

なにかの json が無効な文字列になっていることがわかりました。 .vscode 以下のファイルでエラーが発生しているので、このファイルを見てみることにしました。

~/.vscode/extensions/ms-python.python-2024.20.0-linux-x64/python_files/python_server.py
if __name__ == "__main__":
    while not STDIN.closed:
        try:
            headers = get_headers()
            content_length = int(headers.get("Content-Length", 0))

            if content_length:
                request_text = STDIN.read(content_length)
                # ★ここでエラー発生
                request_json = json.loads(request_text)
                if request_json["method"] == "execute":
                    execute(request_json, USER_GLOBALS)
                if request_json["method"] == "check_valid_command":
                    check_valid_command(request_json)
                elif request_json["method"] == "exit":
                    sys.exit(0)

        except Exception:  # noqa: PERF203
            print_log(traceback.format_exc())

https://github.com/microsoft/vscode-python/blob/2fed954a350142219000ccd2f1579bcdaa9c6e95/python_files/python_server.py#L171-L188

見たところ、 stdin/stdout を使って呼び出し元プログラムと対話していることがわかったので、 ファイルに json の内容を書き出してみることにしました。 (真っ当に pdb を接続するとかも手段かもしれませんが、 拡張機能内から起動されるプログラムであり途中でアタッチするのが難しそうだったのでファイルに吐き出しました)

~/.vscode/extensions/ms-python.python-2024.20.0-linux-x64/python_files/python_server.py
if __name__ == "__main__":
    while not STDIN.closed:
        try:
            headers = get_headers()
            content_length = int(headers.get("Content-Length", 0))

            if content_length:
                request_text = STDIN.read(content_length)
                with open("/tmp/content.txt", mode="w") as cf, open("/tmp/header.txt", mode="w") as hf:
                    cf.write(request_text)
                    hf.write(str(headers))

                request_json = json.loads(request_text)
                if request_json["method"] == "execute":
                    execute(request_json, USER_GLOBALS)
                if request_json["method"] == "check_valid_command":
                    check_valid_command(request_json)
                elif request_json["method"] == "exit":
                    sys.exit(0)

        except Exception:  # noqa: PERF203
            print_log(traceback.format_exc())

以下の様な結果になりました。コンテンツの末尾に変な文字列 ("Co") が入っており、また Content-Length = 82 はどうやら「あ」 を 3 文字換算したもの (utf-8 での表現) ですが、 コンテンツ読む際には 1 文字換算になっているために本来であれば次のヘッダである "Co" が読まれてしまっていることがわかりました。

/tmp/header.txt
{'Content-Length': '82'}
/tmp/content.txt
{"jsonrpc":"2.0","id":0,"method":"execute","params":["print(\"hello2\") # あ\n"]}Co

バグの修正を作る

以上の調査から python 側で sys.stdin.read はどうやら「あ」を1文字換算にしていますが、呼び出し元はそうではないことがわかりました。

https://docs.python.org/3/library/sys.html を読んでみたところ、以下の記載を見つけました。sys.stdio.read に与える引数は utf-8 の文字列長ですが、 sys.stdio.buffer.read を使うことで bytes として読むことができます。

Note To write or read binary data from/to the standard streams, use the underlying binary buffer object. For example, to write bytes to stdout, use sys.stdout.buffer.write(b'abc').
However, if you are writing a library (and do not control in which context its code will be executed), be aware that the standard streams may be replaced with file-like objects like io.StringIO which do not support the buffer attribute.

ということで、以下のように変更することでもともと見つけていたバグは修正できそうでした。

~/.vscode/extensions/ms-python.python-2024.20.0-linux-x64/python_files/python_server.py
if __name__ == "__main__":
    while not STDIN.closed:
        try:
            headers = get_headers()
            content_length = int(headers.get("Content-Length", 0))

            if content_length:
                //bytes で読んでから str に変換
                request_text = STDIN.buffer.read(content_length).decode()
                request_json = json.loads(request_text)
                if request_json["method"] == "execute":
                    execute(request_json, USER_GLOBALS)
                if request_json["method"] == "check_valid_command":
                    check_valid_command(request_json)
                elif request_json["method"] == "exit":
                    sys.exit(0)

        except Exception:  # noqa: PERF203
            print_log(traceback.format_exc())

ちなみに python 側を修正するのではなく、 呼び出し元を修正するというのもありうると思います。ただ、 一般に "Content-Length" はバイト単位であるため、一旦 PR では python 側を修正することにしました。

なお、 python_server.py には他にも Content-Length の扱いが不正なところがあるため、そこも合わせて修正が必要と気づきました。 以下が例で、 len(msg) は utf-8 文字列の長さですが、それを Content-Length として渡してしまっています。

~/.vscode/extensions/ms-python.python-2024.20.0-linux-x64/python_files/python_server.py
def _send_message(msg: str):
    length_msg = len(msg)
    STDOUT.buffer.write(f"Content-Length: {length_msg}\r\n\r\n{msg}".encode())
    STDOUT.buffer.flush()

というわけでいくつかの修正箇所がわかったので issue と PR を作成することにしました。

注意しなければならない点として、 stdin.read*stdin.buffer.read* を同時に使ってはならないというものがありました。これは stdin.read が内部的にバッファリングされているためです。

Upstream での修正

まず https://github.com/microsoft/vscode-python/wiki を読んで PR の作成や Issue の作成方法を学び、以下の PR と Issue を作成しました。

https://github.com/microsoft/vscode-python/issues/24479
https://github.com/microsoft/vscode-python/pull/24480

その後レビューが進み、無事マージされました 🚀

ちなみに同様のバグを後日報告された方もいました。こちらも修正されているはずです。
https://github.com/microsoft/vscode-python/issues/24481

Visual Studio Code の Python 拡張をインストールした人の手元で自分のコードが活躍すると思うとワクワクしますね!

Discussion