pip install中に他所からデータファイルを取得する

概要

そんなことをしても良いのか分からないが，パッケージに必要なデータを PyPI サーバ以外から取得してインストールしたい．
setup.pyに記載できるデータファイル関連の引数，package_data と data_files はどちらもソースコードと共に配布される
ファイルしか対応していないと思われる．
そこで，pip install中に呼び出される install_data コマンドをフックして，必要なファイルを別途ダウンロードするようにした．

コマンドのフック

コマンドのフックというか置き換えは，setuptools.setup関数にキーワード引数 cmdclass として，
置き換えるコマンド名をキーに，実装を値にした辞書を渡すことで行える．(参考: 新しいコマンドの統合)

今回置き換える install_data コマンドの本来のソースは，ここにあるので，
参考にしながらカスタムの install_data コマンドを作成した．

import distutils.command.install_data
from os import path
import site
import sys
import urllib


class CustomInstallData(distutils.command.install_data.install_data):

    def run(self):
        # self.data_files に setup 関数の data_files 引数に渡した値が入る．
        # これは，(データファイルの保存先，データファイルパスのリスト) タプルのリスト担っているので，個別に処理する．
        for f in self.data_files:
            if not isinstance(f, tuple):
                continue
            
            for i, u in enumerate(f[1]):
                # データファイルのパスが URL でなければ何もしない．
                if not u.startswith("http"):
                    continue

                # 対象のデータファイルが既にローカルにあれば，それを再利用する．
                base = path.basename(u)
                f[1][i] = path.join(sys.prefix, f[0], base)
                if not path.exists(f[1][i]):
                    f[1][i] = path.join(sys.prefix, "local", f[0], base)
                if not path.exists(f[1][i]):
                    f[1][i] = path.join(site.getuserbase(), f[0], base)
                if not path.exists(f[1][i]):
                    # ローカルに見つからない場合，ダウンロードする．
                    f[1][i] = urllib.urlretrieve(u, base)[0]

        # その他の処理は，元々のコマンドに移譲する．
        return distutils.command.install_data.install_data.run(self)

基本的な処理は，元々の install_data コマンドと同じにするため，
distutils.command.install_data.install_data を継承したクラスを作成する．
親クラスは，distutils.command.install_data ではなく distutils.command.install_data.install_dataであることに注意．

run メソッドがコマンドの本体なので，このメソッドをオーバーライドする．
self.data_files には， setup 関数の data_files 引数に渡す (directory, files) ペアのリストが格納されている．
directory はデータファイルの保存先で， files はインストールするデータファイルパスのリストである．

self.data_files から個々の要素を取り出して，またその中から個々のデータファイルパスを取り出して処理する．
データファイルパスが URL でなければ通常の処理で良いため，それらに対しては特に何もしない．

次に，対象のデータファイルが既にインストール済みであれば，無駄なダウンロードを避けるために，それを再利用するようにしたい．
データファイルの保存先は，マニュアルによると，

directory が相対パスの場合, インストールプレフィクス (installation prefix、 pure Python パッケージなら sys.prefix, 拡張モジュールの入ったパッケージなら sys.exec_prefix) からの相対パスと解釈されます.

とあるので， 今インストールしようとしているファイル名が filename の場合，

• os.path.join(sys.prefix, directory, filename)
• os.path.join(sys.exec_prefix, directory, filename)

を調べれば良い．ところが，手元の Ubuntu では，上記のいずれでもなくos.path.join(sys.prefix, "local", directory, filename)と一段階掘ったところに保存されていた．
そこで，上記のコードでは， local 付きのパスも調べるようにしている．

また，pip install --user とパッケージをユーザディレクトリにインストールしているケースも考えられる．
この場合，パッケージは site.USER_BASE 以下に保存されるため，
上記のコードでは，site.getuserbase()を使ってsite.USER_BASEの値を取得し，こちらのファイルの存在も調べている．

最後に，ローカルに対象のデータファイルが見つからなかった場合， urllib.urlretrieve を使って取得している．

残りの処理は，元々の install_data コマンドと同じで良いため，親クラスのメソッドを呼び出している．
なお，distutils.command.install_data.install_data は旧クラスのため，distutils.command.install_data.install_data.run(self)によって親クラスのメソッドを呼び出す．

setup 関数

上記のカスタムコマンドを使った setup 関数の呼び出しは次の通り（関係ない項目は省いてある）．

setup(
    data_files=[(
        "rgmining/data",
        ["http://times.cs.uiuc.edu/~wang296/Data/LARA/TripAdvisor/TripAdvisorJson.tar.bz2"]
    )],
    cmdclass={
        "install_data": CustomInstallData
    },
    ... # その他の項目は省略
)

cmdclass 引数に， install_data コマンドの実装クラスとして CustomInstallData を渡している．
これで，data_files に URL を渡すことができるようになる．

まとめ

巨大なファイルをPyPI サーバへ登録することを避けるため， pip install 時に必要なファイルを別サーバからダウンロードする方法をまとめた．
ただし，取得したファイルのハッシュチェックもしていないので，セキュリティ上あまり良くない．
そもそも，こんな面倒なことをしなくても別の方法があるような気がするので，ご存知の方は教えて下さい．