Python实现类型安全可复用大顶堆优先级队列的相关问题咨询

目标 #

• 定义一个基于大顶堆实现的优先级队列
• 利用Python动态类型、仅支持类型提示的特性，在尽可能保障类型安全的同时，实现灵活可复用的队列

我发现网上的多数实现方案，默认大顶堆优先级队列只存储int、float类型的元素。但如果我们需要存储按年龄排序的Person对象、按金额排序的交易（Transactions）对象呢？现有多数实现如果遇到用户插入任意自定义类型，运行时就会报错。

预期目标 #

1. 允许用户将该优先级队列实现复用于其他数据类型
2. 错误使用时快速失败：比如插入队列无法识别比较方式的类实例时及时抛出错误

实现优先级队列需要通过相等/比较运算符对比对象，这里我们假设仅需要用到>、>=两个运算符。
查阅资料后我了解到，自定义类可以通过实现以下魔法方法支持比较操作，能大幅提升灵活性：

==  __eq__
!=  __ne__
<   __lt__
<=  __le__
>   __gt__
>=  __ge__

请问我是否可以在构造函数中检查类型是否实现了所需的比较方法，缺失则抛出异常？如果该思路有问题，还有哪些可行的实现方案？

基础框架代码 #

from typing import TypeVar, Generic, List

# 定义泛型类支持的类型，理想情况下仅接受符合接口的类（即定义了各类比较所需方法的类）
T = TypeVar("T", int, float)


class MaxHeapPriorityQueue(Generic[T]):
    def __init__(self):
        # 此处需要检查类型T是否实现了__gt__、__ge__等必要方法，如果需要的方法有缺失则抛出异常
        self._heap: List[T] = []
        self._insert_pointer: int = 0

    def insert(self, value: T) -> None:
        # 待实现

    def delete_max(self) -> T:
        # 待实现

    def __trickle_up(self, node_index: int) -> None:
        parent_index = self.__calculate_parent_node_index(node_index)

        # 如果用户在堆中存储了非数值类型，这里的元素对比可能失败或出现逻辑bug
        while node_index > 1 and self._heap[node_index] > self._heap[parent_index]:
            self.__exchange(node_index, parent_index)
            node_index = parent_index
            parent_index = self.__calculate_parent_node_index(node_index)

    @staticmethod
    def __calculate_parent_node_index(child_node_index: int) -> int:
        return child_node_index // 2

    def __exchange(self, node_index_1: int, node_index_2: int) -> None:
        # 待实现

使用Protocol的更新版本 #

目前mypy检查可以正常生效，但typing模块不会主动抛出异常，if not isinstance(T, SupportsComparison): raise TypeError('can not instantiate with that type')这段代码的判断分支始终不会触发。

泛型数据结构代码 #

from typing import TypeVar, Generic, List, Protocol, runtime_checkable


@runtime_checkable
class SupportsComparison(Protocol):

    def __lt__(self, other) -> bool: ...

    def __le__(self, other) -> bool: ...

    def __eq__(self, other) -> bool: ...

    def __ne__(self, other) -> bool: ...

    def __ge__(self, other) -> bool: ...

    def __gt__(self, other) -> bool: ...


T = TypeVar("T", bound=SupportsComparison)


class MaxHeapPriorityQueue(Generic[T]):
    def __init__(self):
        if not isinstance(T, SupportsComparison):
            raise TypeError('该类型无法实例化优先级队列')

        self._heap: List[T] = []
        # 插入指针用于保证元素插入后始终形成完全二叉树，指向最新插入的元素位置
        self._insert_pointer: int = 0

    def insert(self, value: T) -> None:
        # 插入前先递增指针，因为指针初始不是指向下一个元素的插入位置，而是比插入位置小1
        # 比如堆中已有1个元素时指针指向1，堆为空时指针指向0
        self._insert_pointer += 1
        self._heap.insert(self._insert_pointer, value)
        self.__trickle_up(self._insert_pointer)

    def delete_max(self) -> T:
        if self._insert_pointer == 0:
            raise Exception("优先级队列为空，无法执行删除操作")
        return self._heap[1]
        # 待实现剩余逻辑

    def __trickle_up(self, node_index: int) -> None:
        parent_index = self.__calculate_parent_node_index(node_index)

        # 如果已经上浮到二叉树根节点，或者当前上浮节点小于父节点，则停止上浮
        while node_index > 1 and self._heap[node_index] > self._heap[parent_index]:
            self.__exchange(node_index, parent_index)
            node_index = parent_index
            parent_index = self.__calculate_parent_node_index(node_index)

    @staticmethod
    def __calculate_parent_node_index(child_node_index: int) -> int:
        return child_node_index // 2

    def __exchange(self, node_index_1: int, node_index_2: int) -> None:
        temp = self._heap[node_index_1]
        self._heap[node_index_1] = self._heap[node_index_2]
        self._heap[node_index_2] = temp

测试用自定义类代码 #

class Person:

    def __init__(self, name, age):
        self.name = name
        self.age = age

    def __lt__(self, other) -> bool:
        return True  # 待实现逻辑，此处仅为测试

    def __le__(self, other) -> bool:
        return True  # 待实现逻辑，此处仅为测试

    def __eq__(self, other) -> bool:
        return True  # 待实现逻辑，此处仅为测试

    def __ne__(self, other) -> bool:
        return True  # 待实现逻辑，此处仅为测试

    def __gt__(self, other) -> bool:
        return True # 待实现逻辑，此处仅为测试

    def __ge__(self, other) -> bool:
        return True  # 待实现逻辑，此处仅为测试


class Animal:

    def __init__(self, breed):
        self.breed = breed

实例化测试代码 #

if __name__ == '__main__':
    max_pq = MaxHeapPriorityQueue[Person]()  # 可通过mypy检查
    max_pq2 = MaxHeapPriorityQueue[Animal]()  # 无法通过mypy检查

内容的提问来源于stack exchange，提问作者MPC