ネコと和解せよ

技術的なあれこれの備忘録のつもり

PyMySQLで取得したレコードのカラムにドットアクセスしたい

tl;dr

  • record["column_name"]でなく、record.column_nameでアクセスしたい
  • DictCursorを参考にDataclassCursorを作る
  • DataclassCursorでdictではなくdataclassを動的に作成する
  • 動的に定義したdataclassのインスタンスをレコードとして返す

DataclassCursor実装例

github.com

今回の流れ

こういう感じのコードを書くとき、 record["column_name"]じゃなくてrecord.column_nameで カラムの情報にアクセスしたくなった。

import pymysql

connection = pymysql.connect(host="", port=, user="", passwd="", db="")

with connection:
    with connection.cursor(pymysql.cursor.DictCursor) as cursor:
        cursor.execute("SELECT * FROM TBL_NAME")
        record = cursor.fetchone()
        record["column_name"]

方法をいろいろ考えてみた結果、dataclassを返すcursorを実装してみることにした。

こんな感じの流れで実装を進める

  1. dataclassを動的に作る

  2. PyMySQLのDictCursorのソースを読む

  3. DictCursorを参考にDataclassCursorを実装

dataclassを動的に作る

dataclassは Python3.7から追加された機能で次のようにメンバ変数を簡素に定義できる

import dataclasses

@dataclasses.dataclass
class User:
    id: int
    name: str
    email: str

user = User(id=0, name="user-name", email="foo@bar.buzz")
print(user.name)

docs.python.org

またこれはdataclasses.make_dataclassを利用することで次のように書くこともできる

import dataclasses

User = dataclasses.make_dataclass(
    'User',
    [('id', int), ('name', str), ('email', str)]
)

docs.python.org

今回はPyMySQLで取得したレコードの情報からmake_dataclassを利用して動的にdataclassを作成する

PyMySQLのDictCursorのソースを読む

DictCursorの実装をみると次のようになっており、処理本体はDictCursorMixinの方だとわかる。

class DictCursorMixin:
    # You can override this to use OrderedDict or other dict-like types.
    dict_type = dict

    def _do_get_result(self):
        super(DictCursorMixin, self)._do_get_result()
        fields = []
        if self.description:
            for f in self._result.fields:
                name = f.name
                if name in fields:
                    name = f.table_name + "." + name
                fields.append(name)
            self._fields = fields

        if fields and self._rows:
            self._rows = [self._conv_row(r) for r in self._rows]

    def _conv_row(self, row):
        if row is None:
            return None
        return self.dict_type(zip(self._fields, row))


class DictCursor(DictCursorMixin, Cursor):
    """A cursor which returns results as a dictionary"""

github.com

dict型に変換する処理の本体はこの部分

    def _conv_row(self, row):
        if row is None:
            return None
        return self.dict_type(zip(self._fields, row))

dict_typedict型がセットされており、_conv_rowでDBから取得したレコード情報をdict型に変換していることが解る。 self._fieldsカラム名のリストで、rowには取得した各カラムの値が入っている。

self.dict_type(zip(self._fields, row))つまりdict(zip(self._fields, row))カラム名をkey, カラムの値をvalueとする辞書配列で、これを元にdataclassを生成すれば良さそうなことが解った

DataclassCursorを実装

次のように辞書配列を受け取り、dataclassを動的に生成し、 そのdataclassのインスタンスを生成する関数を実装する

import dataclasses

def to_dataclass(data: dict, frozen: bool = False):
    fields = [(key, type(val)) for key, val in data.items()]

    data_cls = dataclasses.make_dataclass(
        'Record',
        fields,
        frozen=frozen
    )

    return data_cls(**data)

そしてDictCursorMixinを継承したDataclassCursorMixinで、 _conv_rowを 上記で実装したto_dataclassを利用した処理に書き換える

import pymysql

class DataclassCursorMixin(pymysql.cursors.DictCursorMixin):
    def _conv_row(self, row):
        if row is None:
            return None
        return to_dataclass(dict(zip(self._fields, row)), frozen=False)


class DataclassCursor(DataclassCursorMixin, pymysql.cursors.Cursor):
    """DataclassCursor"""

使い方

import pymysql

connection = pymysql.connect(host="", port=, user="", passwd="", db="")

with connection:
    with connection.cursor(DataclassCursor) as cursor:
        cursor.execute("SELECT * FROM TBL_NAME")
        record = cursor.fetchone()
        record.column_name

今回実装したものは下記リポジトリに置いてます

github.com

c.f

github.com

docs.python.org

github.com