SSHの環境を検知して自動でmatplotlibのGUIをブラウザで表示する

Abstract

SSHで接続している時は、matplotlibを利用してプロットしていると通常はGUIが使えません。
Matplotlibでは、そのような状況ではバックエンドを変更することで対処することができます。
しかしながら、次のようなケースでは、面倒なことになります。

• 普段は研究室(職場)のPCで作業をしているが、たまにSSHで自宅から接続して、自宅のPCで作業を行う

普段使う時と自宅から使う時で毎回matplotlibのバックエンドを変更するのは面倒です。
本記事では、SSH接続中かどうかを判定して、自動でバックエンドを変更する仕組みを紹介します。

結論

import os
import atexit
import matplotlib as mpl

def is_under_ssh_connection():
    # The environment variable `SSH_CONNECTION` exists only in the SSH session.
    return 'SSH_CONNECTION' in os.environ.keys()

def use_WebAgg(port = 8000, port_retries = 50):
    """use WebAgg for matplotlib backend.
    """
    current_backend = mpl.get_backend()
    current_webagg_configs = {
        'port': mpl.rcParams['webagg.port'],
        'port_retries': mpl.rcParams['webagg.port_retries'],
        'open_in_browser': mpl.rcParams['webagg.open_in_browser'],
    }
    def reset():
        mpl.use(current_backend)
        mpl.rc('webagg', **current_webagg_configs)
    
    mpl.use('WebAgg')
    mpl.rc('webagg', **{
        'port': port,
        'port_retries': port_retries,
        'open_in_browser': False
    })
    atexit.register(reset)


if is_under_ssh_connection():
    use_WebAgg(port = 8000, port_retries = 50)

普段実行するmatplotlibもファイルで、上記コードを貼り付ければオッケーです。
VSCodeで実行していると、自動でPort Forwadingが行われるので、指示通りブラウザで開くもヨシ。VSCodeのタブで開くもヨシ。

解説

SSH接続のセッションかどうかを判定する

SSH_CONNECTIONという環境変数を確認することで、セッションがSSH接続かどうかを判定できます。

Pythonでは、環境変数はosモジュールを利用すればよいので、次のように書けば判定が可能です。

def is_under_ssh_connection():
    return 'SSH_CONNECTION' in os.environ.keys()

WebAggとは

matplotlibはバックエンドを切り替えることで、様々な環境で作成したグラフを開くことができます。
多くの場合、通常はバックエンドは Qt4Agg などです。
これは描画する処理を何にするかや、GUIの表示方法をどうするかを定義していて、普段plt.show()するとGUIが開くのはこのバックエンドがGUIのウィンドウを作って開くものに設定されているからです。

WebAggでは、ブラウザでグラフを表示することを可能にするバックエンドの一つです。
plt.show()すると、設定されたポートでサーバーを立てて表示します。

その他にもたくさんありますので、公式ドキュメントに目を通してみてください。

atexit.registerとは

プログラムが終了する時に呼ばれる関数を定義できる標準モジュールを利用した実装です。
mpl.use('WebAgg')はグローバルに設定を変更するので、プログラムが終了したら実行前に戻るようにしています。
直前のものではなく、既定のものに戻す場合はdef reset()の部分を次のようにしてみてください。

def reset():
    mpl.use(mpl.rcParamsDefault['backend'])
    ...