Pythonのdataclassを使ってみた

Pythonの標準ライブラリにdataclassというのがあるの見つけたので使ってみました。
参考: dataclasses — データクラス — Python 3.10.6 ドキュメント

名前から、オリジナルのデータ型を定義するためのモジュールなのかなとも思ったのですが実際は少し違いそうです。もちろん、オリジナルのデータ型を定義するためにも使えるのですが、その実態は、クラスに対して__init__()や__repr__()といった特殊メソッドを自動的に生成してくれるデコレーターという解釈が正確のようです。

お試しに、証券コードと会社名と説明を属性として持ったCompanyクラスを作ってみましょう。

import dataclasses


@dataclasses.dataclass
class Company:
    code: int
    name: str
    description: str = None


# __init__() メソッドが自動的に定義されているためこれでインスタンスを作成できる
toyota = Company(code=7203, name="トヨタ自動車", description="自動車メーカー")
suzuki = Company(code=7269, name="スズキ", description="自動車メーカー")


# __eq__() メソッドが自動的に定義されているため、比較ができる
toyota == suzuki
# False

これは属性がたった3個だけで、メソッドも持ってないようなクラスなのですが、__init__()が入らないってだけでものすごくシンプルに描けるようになりましたね。

また、比較用のメソッドを自分で作らなくても、各属性が全て一致しているかどうかを基準に一致不一致を判定してくれるのも便利です。属性が3個だけだとそこまでありがたみがないですが、もっと大規模なクラスで、全属性の一致を判定するのは無駄にコードが長くなりますから。

int とか str と書いて型ヒントをつけられたりするのも今風な感じがします。ただ、この型ヒントはどうやらただのアノテーションで、代入する値に対する強制力などはないようです。
codeを整数、nameを文字列としていますがそうでない値も入ります。

dummy_company = Company(code="文字列", name=1234)
print(dummy_company)
# Company(code='文字列', name=1234, description=None)

この記事の冒頭で書いていますが、このdataclassはclassに対するデコレーターなので、ただのオリジナルデータ型ではなく、普通にメソッド等を持っているクラスを作成することもできます。その場合__init__()などを自分で書かなくて良くなるので、特に凝った__init__()が不要な場合はバンバン使って良さそうです。例えば、二つの値を持ち、合計値を返せるクラスは次のようになります。

@dataclasses.dataclass
class two_number:
    a: int
    b: int

    def sum(self):
        return self.a + self.b


tn = two_number(5, 8)
print(tn.sum())
# 13

自分が実装したいメソッドの部分に専念できるのはいいですね。

このdataclassのデコレーターですが、デコレーター自体も引数を取って、色々設定することができます。詳細はドキュメントに譲りますが、例えばinit=Falseやrepr=False, eq=Falseを指定すると、デフォルトで生成されると言っていた__init__()や__repr__()、__eq__()などが生成されなくなります。自分で実装したいものがあったらそれだけ自分で実装するようにしましょう。

frozen (デフォルトはFalse) を Trueに指定すると、値への代入が禁止されます。これのメリットしては辞書のキーとして使えるようになることでしょうか。ちょっとやってみます。
１個目の例は上で作ったCompanyなので、frozenはFalseです。その次がfrozen=True。

# frozen = False だと属性に値を代入できる
toyota.description  = "日本の自動車メーカー"
print(toyota)


# frozen=Trueを指定してみる
@dataclasses.dataclass(frozen=True)
class Frozen_Company:
    code: int
    name: str
    description: str = None


f_toyota = Frozen_Company(code=7203, name="トヨタ自動車", description="自動車メーカー")

# frozen = True だと属性に値を代入できないため、例外が上がる
try:
    f_toyota.description  = "日本の自動車メーカー"
except Exception as e:
    print(type(e), ":", e)
#  <class 'dataclasses.FrozenInstanceError'> : cannot assign to field 'description'

frozenにするメリットとしては、タプルと同様に辞書のキーにできる、という点があります。

# frozenではない、つまりハッシュ化不可能なので辞書のキーにできない
try:
    {toyota: 1}
except Exception as e:
    print(type(e), ":", e)
# <class 'TypeError'> : unhashable type: 'Company'

# frozen=Trueだとハッシュ化可能なので辞書のキーにできる
{f_toyota: 1}
# {Frozen_Company(code=7203, name='トヨタ自動車', description='自動車メーカー'): 1}

また、order という引数(デフォルトFalse)にTrueを渡すと、__le__()等々の不等号を実装する特殊メソッドたち4種も自動的に生成してくれるようになります。どうも要素を順番に比較して最初に上下がついたもので決まるようです。これも属性が多い時は便利なのではないでしょうか。

@dataclasses.dataclass(order=True)
class two_number:
    a: int
    b: int

tn1 = two_number(12, 5)
tn2 = two_number(6, 20)
tn1 > tn2
# True

さて、以上でdataclass自体の基本的な説明はおしまいです。

あとは偶然気づいた豆知識なのですが、Pandasとの連携について紹介します。
dataclassの配列は、簡単にPandasのDataFrameに変換できます。実質的にdictみたいに振る舞ってくれるようです。

import pandas as pd


df = pd.DataFrame([toyota, suzuki])
print(df)
"""
   code    name description
0  7203  トヨタ自動車  日本の自動車メーカー
1  7269     スズキ     自動車メーカー
"""

便利ですね。

逆に、DataFrameに入った値たちをdataclassで定義したクラスのインスタンスに変換したいな、と思って方法探しました。専用のメソッドなどは見つかってないのですが、lamda関数を使ってこのようにするのが良いでしょう。

df.apply(lambda row: Company(**row), axis=1)
"""
0    Company(code=7203, name='トヨタ自動車', description=...
1    Company(code=7269, name='スズキ', description='自動...
dtype: object
"""

事前にDataFrameの列名と、dataclassの属性名を揃えておく必要はあるのでそこは注意してください。

コメントを残す コメントをキャンセル

コメントを残すコメントをキャンセル