有点意思的 Python 系列二 内置类型增加额外方法

我们知道，如果想给一个自定义类型增加对应的方法，可以直接修改这个类就可以了，但是如果我们这个时候想给内置的类型增加一些自定义方法呢？

比如对于可变类型对象，我们想增加一个 deepcopy 的方法实现深拷贝，类似dd = {'a': [1, 2, 3]}.deepcopy() 达到 dd = copy.deepcopy({'a': [1, 2, 3]}) 的效果，显而易见的是直接在一个 dict 对象上 .deepcopy() 是更优雅的。但是事与愿违，我们没有办法通过常规手段给 dict 增加 deepcopy() 方法。

难道就真的没有办法吗？下面这段代码可以优雅的实现。

import ctypes


class PyType(ctypes.Structure):
    pass


class PyObject(ctypes.Structure):
    Py_ssize_t = (
        ctypes.c_int64 if ctypes.sizeof(ctypes.c_void_p) == 8 else ctypes.c_int32
    )
    _fields_ = [
        ("ob_refcnt", Py_ssize_t),
        ("ob_type", ctypes.POINTER(PyType)),
    ]


class PyTypeObject(PyObject):
    _fields_ = [
        ("dict", ctypes.POINTER(PyObject))
    ]


def inject(class_, method, force=False):
    def _(function):
        name_, dict_ = class_.__name__, class_.__dict__
        proxy_dict = PyTypeObject.from_address(id(dict_))
        namespace = {}
        ctypes.pythonapi.PyDict_SetItem(
            ctypes.py_object(namespace),
            ctypes.py_object(name_),
            proxy_dict.dict
        )
        if not force and namespace.get(name_, {}).get(method, None):
            raise RuntimeError(f"已存在方法 {class_.__name__}.{method}()")
        namespace[name_][method] = function

    return _

而使用使用方法也很简单，比如上面的给 dict 添加一个 deepcopy() 实现字典的深拷贝

import copy

@inject(dict, 'deepcopy')
def deepcopy(d):
    return copy.deepcopy(d)

# 验证一下
origin_dict = {"goods": ["apple", "orange"]}

print(f"初始字典: {origin_dict}") # {'goods': ['apple', 'orange']}

copy_dict = origin_dict.copy() # 自带的 copy() 浅拷贝
deepcopy_dict = origin_dict.deepcopy() # 添加的 deepcopy() 深拷贝

origin_dict["goods"].append("banana")

print(f"初始字典变更: {origin_dict}") # {'goods': ['apple', 'orange', 'banana']}
print(f".copy() 结果: {copy_dict}") # {'goods': ['apple', 'orange', 'banana']}
print(f".deepcopy() 结果: {deepcopy_dict}") # {'goods': ['apple', 'orange']}

再或者给 list 添加一个 average() 方法计算平均数

@inject(list, 'average')
def average(l):
    return sum(l) / len(l)

score = [95.0, 89.5, 77.0, 91.0]

print(score.average())  # 88.125

再或者给字符串添加一个 json() 方法，可以直接通过 str.json() 将该字符串格式化为 json 对象（当然前提是这个字符串是可以被反序列化为 json 对象）

from json import loads

@inject(str, 'json')
def json(s):
    return loads(s)

info = '{"first_name": "Michael", "last_name": "Rodgers", "department": "Marketing"}'

print(info.json()) # {'first_name': 'Michael', 'last_name': 'Rodgers', 'department': 'Marketing'}

同样的，比如给 int 类型添加 add(number) 方法

@inject(int, 'add')
def add(i, number):
    return i + number

munber = 5

print(number.add(3)) # 8

# 还可以进行链式调用
print(munber.add(3).add(7).add(-1)) # 14

当然除了内置类型，也可以修补自定义类型

class Number(object):
    def __init__(self, n):
        self.number = n

@inject(Number, 'sub')
def sub(n, num):
    return Number(n.number - num)

number = Number(10)
print(number.sub(3).sub(5).number)  # 2