Python: pandas-should というパッケージを作ってみた

pandas を使ってデータ分析などをしていると、自分が意図した通りのデータになっているか、たまに確認することになると思う。 確認する方法としてはグラフにプロットしてみたり、あるいは assert 文を使って shape などを確認することが考えられる。

今回紹介する pandas-should は後者の「assert 文を使った内容の確認」を、なるべく簡単に分かりやすく記述するために作ってみた。

github.com

使った環境は次の通り。 なお、パッケージ自体は Python 3.6 以降で動作する。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.14.5
BuildVersion:   18F132
$ python -V
Python 3.7.4

インストール

インストールは pip からできる。

$ pip install pandas-should

使い方

パッケージをインストールできたら、とりあえず Python のインタプリタを起動しておく。

$ python

あとは pandas_should をインポートするだけ。

>>> import pandas_should

インポートすると pandas の DataFrame と Series のインスタンスに should というアトリビュートがひそかに生えてくる。

>>> import pandas as pd
>>> df = pd.DataFrame([1, 2, 3], columns=['id'])
>>> df.should
<pandas_should.dataframe.ShouldDataFrameAccessor object at 0x1083d6160>
>>> s = pd.Series([1, 2, 3])
>>> s.should
<pandas_should.series.ShouldSeriesAccessor object at 0x1196a36a0>

この should 経由で色々とできて、例えば行数が一致することを確認したいなら have_length() を使う。

>>> df.should.have_length(3)
True

基本的にメソッドは真偽値を返すので、アサーションに使うならこうする。

>>> assert df.should.have_length(3)

ここからは使いそうな API を幾つか紹介していく。

DataFrame

まずは DataFrame から。

要素に Null (NaN or NaT) が含まれるか調べたい

普通に書くと、こんな感じになると思う。

>>> not df.isnull().any(axis=None)
True

pandas-should を使うと、こう書ける。

>>> df.should.have_not_null()
True

あるいは Null が含まれることを期待するのであれば、こう。

>>> df.should.have_null()
False

要素のレンジを調べたい

各要素が特定のレンジ (値の範囲) に収まっているか知りたいときは、こう書く。 値の範囲には、指定した最小値と最大値も含まれる。

>>> df.should.fall_within_range(1, 3)
True

下限だけ指定したいときは greater_than() を使う。

>>> df.should.greater_than(0)
True

greater_than() では指定した値は含まれないので、含みたいときは greater_than_or_equal() を使う。

>>> df.should.greater_than_or_equal(1)
True

長いのでエイリアスとして gt() や gte() も使える。

>>> df.should.gt(1)
False
>>> df.should.gte(1)
True

上限についても同様。 こちらもエイリアスとして lt() と lte() が使える。

>>> df.should.less_than(3)
False
>>> df.should.less_than_or_equal(3)
True

形状 (Shape) を調べたい

続いて DataFrame の形状を調べる方法について。

比較対象が必要なので新たに DataFrame を用意しておく。

>>> data1 = [
...     ('apple', 98, True),
...     ('banana', 128, True),
... ]
>>> df1 = pd.DataFrame(data1, columns=['name', 'price', 'fruit'])
>>> data2 = [
...     ('carrot', 198, False),
...     ('dates', 498, True),
... ]
>>> df2 = pd.DataFrame(data2, columns=['name', 'price', 'fruit'])

同じ行数や列数であることを確認したいときは have_same_length() や have_same_width() を使う。

>>> df1.should.have_same_length(df2)
True
>>> df1.should.have_same_width(df2)
True

前述した通り、整数で指定したいときは have_width() や have_length() が使える。

>>> df1.should.have_width(2)
True
>>> df1.should.have_length(2)
True

ちなみに have_same_*() は複数の DataFrame との比較もできる。

>>> data3 = [
...     ('eggplant', 128, False),
... ]
>>> df3 = pd.DataFrame(data3, columns=['name', 'price', 'fruit'])
>>> data4 = [
...     ('fig', 298, True),
... ]
>>> df4 = pd.DataFrame(data4, columns=['name', 'price', 'fruit'])

例えば二つの DataFrame の行数を加算したものと同じになるか調べたいときは以下のようにする。 具体的なユースケースとしては、結合前と結合後の DataFrame の行数が一致しているか調べるときとか。

>>> df1.should.have_same_length(df3, df4)
True

行数についても同様。

>>> df1.should.have_same_width(df3, df4)
True

行と列を別々に比較するのがめんどいときは be_shaped_like() で一気に比較できる。

>>> df1.should.be_shaped_like(df2)
True

このメソッドにはタプルとか整数も渡せる。

>>> df1.should.be_shaped_like(df2.shape)  # tuple
True
>>> df1.should.be_shaped_like(df2.shape[0], df2.shape[1])  # int, int
True

Series

続いては Series について。

要素に Null (NaN or NaT) が含まれるか調べたい

要素に Null が含まれるか調べたいときは DataFrame と同じやり方が使える。

>>> s.should.have_not_null()
True
>>> s.should.have_null()
False

要素のレンジを調べたい

Series に関しても DataFrame と同じように、要素のレンジ (値の範囲) を調べられる。 追加で説明することは特にないかな。

>>> s.should.fall_within_range(1, 3)
True
>>> s.should.gt(1)
False
>>> s.should.gte(1)
True
>>> s.should.lt(3)
False
>>> s.should.lte(3)
True

形状 (Shape) を調べたい

Series に関しては列数という概念がないけど、次のように行数に関しては DataFrame と同じやり方が使える。

>>> s2 = pd.Series([4, 5, 6])
>>> s.should.have_same_length(s2)
True
>>> s.should.have_length(3)
True

そんなかんじで。 こういう API があると便利で欲しいみたいなのがあれば教えてほしい。

スマートPythonプログラミング: Pythonのより良い書き方を学ぶ