IPythonでimport文を省略したいときの最適解 #Python

はじめに

ターミナル上でIPythonを実行したりJupyter QtConsole使ったりというときに、毎回import matplotlib.pyplot as pltと打ったり、NameError: name 'plt' is not definedでイラついたりを避けたいですよね。

よく見る解決策

IPythonは起動時にユーザーディレクトリからたどって /.ipython/profile_default/startupディレクトリ直下にあるpyファイルを全部（多分）実行します。この中にたとえば00-startup.pyファイルを

# 00-startup.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

で保存しておくと、IPythonの起動後すぐにnpとか使えます

欠点

起動時にimportしているので、当然起動が遅くなります。特にpandasやmatplotlibはかなり重いので、起動がだいぶ遅れます。

できれば、pltを使わないときはimportせず、使うときは使うときでわざわざimportせずにどうにかしてほしいです。

最適解！

00-startup.pyに次のLazyImporterを定義します（少しずつ解説します）。

## 00-startup.py

class LazyImporter:
    def __init__(self, name: str, alias: str = None):
        self._name = name
        self._alias = alias or name
    
    def __repr__(self):
        return repr(self._mod)
    
    def __str__(self):
        return str(self._mod)
    
    def __getattr__(self, name):
        return getattr(self._mod, name)

    @property
    def _mod(self):
        from IPython import get_ipython
        import importlib
        mod = importlib.import_module(self._name)
        shell = get_ipython()
        shell.push({self._alias: mod})  # update namespace
        return mod

何がやりたいかというと、np = LazyImporter("numpy", "np")で初期化し、np.arrayとか呼び出したときにnumpyをimportしてグローバルの名前空間のnpを上書きするという作戦です。ポイントは_modの中身です。

class LazyImporter:
    ...
    @property
    def _mod(self):
        from IPython import get_ipython
        import importlib
        mod = importlib.import_module(self._name)  # (1)
        shell = get_ipython()                      # (2)
        shell.push({self._alias: mod})             # (3)
        return mod

各行を説明すると:

import_moduleは引数にモジュールを文字列で渡してimportされたモジュールオブジェクトを返します。self._nameが"numpy"ならmodにnumpyが入り、mod.arrayとかが使える状態です。
get_ipython()は現在起動中のIPythonシェルを返します。ターミナル内でIPythonを起動している場合とJupyter QtConsoleを起動している場合では異なるオブジェクトが返されますが、当然共通のメソッドが実装されています。
pushメソッドはシェルのグローバルな名前空間をdictで更新します。ここで{"np": numpy}を渡せばnpが使えるようになるわけです。

あとは、モジュールが必要になる場面でこの_modプロパティを呼び出すように細工すればよいです。考えられる場面はほぼnp.arrayとかのようにモジュール内の関数やクラスにアクセスするときなので、__getattr__で_modを呼び出し、名前空間を更新しつつモジュールの__getattr__を呼べばよいです。

class LazyImporter:
    ...
    def __getattr__(self, name):
        return getattr(self._mod, name)

一応、>>> npとか>>> print(np)にも対応させるため、__repr__も実装しておいたのが上の例になります。

こうすれば、使う可能性のあるモジュールを

## 00-startup.py

np = LazyImporter('numpy', 'np')
plt = LazyImporter('matplotlib.pyplot', 'plt')
pd = LazyImporter('pandas', 'pd')
sns = LazyImporter('seaborn', 'sns')
da = LazyImporter('dask.array', 'da'))

みたく片っ端から名前空間にぶち込んでも起動は早いままですね！

残る問題点

さすがに変なことをしているので、問題は残ります。例えば、numpyがimportされていない状態でnp.aとまで打ってTabで予測変換をしようとすると落ちます。これが回避できるかはもう少し試してみます。

2022/09/23 追記

import lockのせいでimport_moduleで例外が生じているんだろうなと思ってはいたんですが、try/exceptだけでどうやら落ちるのは回避できるっぽいです。

class LazyImporter:
    ...
    @property
    def _mod(self):
        from IPython import get_ipython
        import importlib
        try:
            mod = importlib.import_module(self._name)
        except Exception:
            return
        shell = get_ipython()
        shell.push({self._alias: mod})
        return mod