Pylance が pd.read_csv() の戻り値を DataFrame として認識してくれない

実行環境
Windows10
Python 3.9.5
Visual Studio Code 1.63.2
Pylance 2022.1.5
pandas 1.3.1

はじめに

python のコードエディタとして Visual Studio Code (以下 VSCode)を使っています。

VSCode には拡張機能として Pylance を入れています。

型チェック (python.analysis.typeCheckingMode) の分析レベルは basic に設定しています。

まずこれを見て下さい。

import pandas as pd

fp = r'data.csv'
df = pd.read_csv(fp)
print(df.head())

csv ファイルを pandas の DataFrame として読み込んで、上から5行を出力するだけのコードです。

もちろん正しく実行できるコードなのですが、エディタ上には警告がでます。

head() に対するこの警告をどうにかしたい というのが今回の趣旨です。

ただ、結論としてはすっきりする解決策にはたどり着けませんでした。

あれこれ調べて考えました、という記事です。

警告の内容

赤い波線の出ている head() をホバーして、警告の内容を見てみます。

Pylance から reportGeneralTypeIssues が出ています。

変数 df の型は pd.DataFrame のつもりなのですが、TextFileReader として認識されているようです。

変数 df にホバーするとこうです。

やはり、TextFileReader になっています。

pandas.read_csv()

csvファイルを読み込んで DataFrame に格納する関数です。

引数が 52 個もあります。

リファレンスには、戻り値は DataFrame か TextParser のどちらかと書いてあります。

Returns: DataFrame or TextParser
A comma-separated values (csv) file is returned as two-dimensional data structure with labeled axes.

この TextParser というのが、pylance の警告にある TextFileReader のことだと思います。

リファレンスページを TextFileReader で検索してみると、2 つの引数（iterator と chunksize）の説明箇所にヒットします。

Return TextFileReader object ... と記載があります。

iterator: bool, default False
Return TextFileReader object for iteration or getting chunks with get_chunk().

chunksize: int, optional
Return TextFileReader object for iteration. See the IO Tools docs for more information on iterator and chunksize.

あまり深入りはしませんでしたが、大きいデータをチャンクに分けて処理したい場合にデータをイテレータとして取得する、という場面で使うようです。

たしかに pd.read_csv() の戻り値が TextFileReader というオブジェクトになり得るということは分かりました。

しかし、iterator 引数のデフォルトは False で、chunksize は optional です。

これらの引数を明示的に使っていないのに、戻り値が TextFileReader 扱いされてしまうのは理不尽です。

Pylance の pd.read_csv() の型定義

コード上の read_csv() をホバーします。

ポップアップに表示されているは pd.read_csv() の Pylance による型の定義文です。

全てを表示できていませんが、実際はこの下にも延々続きます。

この定義文は以下のファイルに記述されているようです。

C:\Users\xxx\.vscode\extensions\ms-python.vscode-pylance-2022.1.5\dist\bundled\stubs\pandas\io\parsers.pyi

もうこのあたりになると、私にとってはムズカシイ話になってきます。

これは型情報を記述したスタブファイルで、この中で @overload デコレータを使って、関数の引数と戻り値の組み合わせが定義されます。

.pyi というのは、型情報を記述したスタブファイルの拡張子です。

parsers.pyi の一部抜粋です。

# -略-
@overload
def read_csv(
    reader: IO,
    sep: str = ...,
    delimiter: Optional[str] = ...,
    # -略-
) -> TextFileReader:
# -略-

parsers.pyi には @overload 付きの read_csv() の定義が 6 つあり、TextFileReader を返すものが 4 つ、DataFrame を返すものが 2 つあります。

そのうち、一番上の 定義文は TextFileReader を返すもので、pd.read_csv() に位置引数でファイルパスのみを指定する場合は、これが適用されるようです。

pd.read_csv() の第一引数

ここまで見てきて気付いたのですが、pd.read_csv() の第一引数の名前が、リファレンス（=ソースコード）とスタブファイルで違います。

pd.read_csv() のソースコードでは、第一引数の名前が filepath_or_buffer で、Pylance のスタブファイルでは reader または filepath です。

試しに、csv ファイルのパスを、位置引数ではなく 3 種類のキーワード引数で指定してみます。

code

import pandas as pd

fp = r'data.csv'

df = pd.read_csv(filepath_or_buffer=fp)
print(df.head())

df = pd.read_csv(reader=fp)
print(df.head())

df = pd.read_csv(filepath=fp)
print(df.head())

filepath 指定で、警告が消えました。

これで解決か、と一瞬思いましたが、これではそもそも実行エラーでした。

filepath_or_buffer 指定のみが正常に動きます。

対策案

以上をふまえて pd.read_csv() の戻り値に出る警告をなんとかするための対策をいくつか考えます。

案1. スタブファイルをカスタムする

Pylance の設定に python.analysis.stubPaths があります。

ここに、カスタムのスタブファイルを追加できるようです。

精通した人はここでなんとかするんでしょうか、、、

私には高いハードルです。

案2. スタブファイルを修正する

read_csv() の型定義が書かれた先述のスタブファイル（parsers.pyi）を、思いのままに修正します。

実は、定義文の戻り値部分を無理やり TextFileReader から DataFrame に変更したら、とりあえず警告は消えました。

ただ、ブラックボックスの中身をいじっているようで怖いです。

また、Pylance のアップデート時に parsers.pyi が上書きされ元に戻ると思われます。

案3. 型チェックをオフにする

Pylance の設定の python.analysis.typeCheckingMode を off にします。

つまり、デフォルトに戻します。

型チェックはしたいので、本末転倒です。

案4. 診断の重大度レベルを下げる

Pylance の設定 python.analysis.diagnosticSeverityOverrides で型チェックの重大度レベルを下げます。

出ている警告は reportGeneralTypeIssues で、今は赤い波線が出ています。

この警告の重大度レベルを warning 以下にします。

重大度レベル
・error (赤い波線)
・warning (黄色い波線)
・information (青い波線)
・none (無効)

settings.json には、例えば以下の記述を追加します。

{
    "python.analysis.diagnosticSeverityOverrides": {
        "reportGeneralTypeIssues" : "warning",
    }
}

型チェックを無効にしたいわけではないし、黄色や青色でもまだ気持ち悪いです。

案5. assert 文（または if 文）で型を絞る

pd.read_csv() の戻り値の型をコード中で、わざわざ判定します。

code

import pandas as pd

fp = r'data.csv'
df = pd.read_csv(fp)

assert isinstance(df, pd.DataFrame)

# または
# if not isinstance(df, pd.DataFrame):
#     raise Exception

print(df.head())

警告は消えますが、変数 df の型判定は以下のようになります。

さらに、クラス定義文の中では def 文をまたぐと警告が復活します。

code

import pandas as pd

class TestDataFrameAssert:
    def __init__(self, fp: str) -> None:
        self.df = pd.read_csv(fp)
        assert isinstance(self.df, pd.DataFrame)

    def print_head(self) -> None:
        print(self.df.head())

案6. 型を cast する

typing モジュールの cast() 関数で pd.Dataframe であることを強制します。

code

from typing import cast
import pandas as pd

fp = r'data.csv'
df = cast(pd.DataFrame, pd.read_csv(fp))

print(df.head())

cast() の中で pd.read_csv() を使います。

望まれないインポート文が cast() のために必要になりますが、変数 df は pd.DataFrame として認識され、警告が消えます。

クラス定義文のなかで def 文をまたいでも、警告がでません。

code

from typing import cast
import pandas as pd

class TestDataFrameCast:
    def __init__(self, fp: str) -> None:
        self.df = cast(pd.DataFrame, pd.read_csv(fp))

    def print_head(self) -> None:
        print(self.df.head())

採用案（=妥協案）

しばらくは、案6 の cast() を使うことにします。

VSCode や Pylance の設定をいじる必要がなく、クラス定義文の中でも効くためです。

対処療法的な方法で腑に落ちないですが、警告を消すという当初の目的は達成できます。

案1 に近いかたちで、スタブファイルを使ったスマートな方法があるのだと思いますが、今回はギブアップです。

さいごに

VSCode の拡張機能 Pylance で、pandas の read_csv() 関数の戻り値が pd.DataFrame 判定にならない原因と対策案を考えました。

いくつか対策案を書きましたが、どの方法もオススメをするものではないので、ご了承ください。

型チェックの仕組みやスタブファイルの作り方については、今後理解を深めていきたいところです。

追記 2022/02/25

本稿投稿時点の pylance のバージョンは 2022.1.5 でしたが、先日 2022.2.4 で改めて確認してみたところ、read_csv() の戻り値が正常に pd.DataFrame と判定されるようになっていました。

警告表示が消えています。

pylance の更新履歴を確認したところ、バージョン 2022.2.0, 2022.2.3, 2022.2.4 で Pandas のスタブについて修正した旨が記載されていました。

これで気持ちよくコーディングできそうです。