控制 Python 类的序列化过程

问题

有的类是不支持在多进程间传递的，如果非要这么做，可能会引发奇怪的现象。比如下面这段代码：

from concurrent.futures import ProcessPoolExecutor, as_completed

from pymysql import connect


class MySqlDatabase(object):

    def __init__(self, host='127.0.0.1', port=3306, user='root', pwd='', db=None):
        self.host = host
        self.port = port
        self.user = user
        self.pwd = pwd
        self.db = db
        self.conn = self.connect()

    def connect(self):
        return connect(host=self.host, port=self.port, user=self.user, password=self.pwd, database=self.db)

    def runquery(self, q):
        with self.conn.cursor() as cur:
            cur.execute(q)
            return cur


def run_in_pool(db, sql):
    print(db.runquery(sql))


def run():
    quires = ['show tables',
              'select * from user']
    db = MySqlDatabase(pwd='1234', db='mysql')
    with ProcessPoolExecutor() as e:
        fs = [e.submit(run_in_pool, db, q) for q in quires]
        for f in as_completed(fs):
            print(f.result())


if __name__ == '__main__':
    run()

这段代码运行后会卡在 for f in as_completed(fs): 这一行，有的平台会抛出异常 TypeError: cannot serialize '_io.BufferedReader' object，有的则什么都不显示，这是因为 pymysql 提供的 Connection 对象是不可序列化的，因此多为多进程间的参数传递会产生异常。

那么，如果让这里的 MySqlDatabase 实例支持在多进程间传递呢？

解决方案

Python 提供了 __getstate__ 和 __setstate__ 方法以支持类进一步控制其实例的封存过程，这会使的实例可以被 pickle 序列化，正确实现这两个方法，实例在多进程间就可以正常传递了。

以上面的问题为例，更改 MySqlDatabase 实现如下：

class MySqlDatabase(object):

	...

    def __getstate__(self):
        state = self.__dict__.copy()
        # 移除不可序列化的属性
	    state.pop('conn')
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        # 重新绑定移除的属性
        self.conn = self.connect()

更新后，再次运行，就会得到预期的结果了。

扩展

• 如果类定义了 __getstate__()，它就会被调用，其返回的对象是被当做实例内容来封存的，否则封存的是实例的 __dict__
• 当解封时，如果类定义了 __setstate__()，就会在已解封状态下调用它。此时不要求实例的 state 对象必须是 dict。没有定义此方法的话，先前封存的 state 对象必须是 dict，且该 dict 内容会在解封时赋给新实例的 __dict__
• 该方法同样适用于 copy.copy 以提供新的拷贝对象