【Python】PriorityQueueで優先度が同じ場合キューに追加した順番で取得できるようにする

はじめに

Pythonのビルドインのqueueモジュールには様々なキュークラスが定義されています．
その中に優先度付きのキュー（PriorityQueue）があります．

優先度が同じ場合，キューに追加した順番でアイテムを取得できるように思えるが，実は違います．
PriorityQueueではヒープキューを使用するのが一般的です．
Pythonでもヒープキュー（heapq）を使って実装しています．
このヒープキューの実装で，キューからポップした際に順番がずれてくるからです（詳しい説明は省きますが要望があれば書きます）．
（ヒープ自体の説明はこちらのサイトが詳しいです）
試したソースが以下です．

from dataclasses import dataclass, field
from typing import Any
import queue


@dataclass(order=True)
class QueueItem:
    """
    キューアイテム
    """
    priority: int
    item: Any = field(compare=False)

def main():
    # キューを用意
    priority_queue = queue.PriorityQueue()
    for index in range(15):
        priority = index % 3
        queue_item = QueueItem(priority=priority, item=index)
        priority_queue.put(queue_item)
        print(f"[input] priority:{priority} value:{index}")

    # 取り出し
    while not priority_queue.empty():
        queue_item = priority_queue.get()
        priority = queue_item.priority
        value = queue_item.item
        print(f"[output] priority:{priority} value:{value}")

if __name__ == "__main__":
    main()

[input] priority:0 value:0
[input] priority:1 value:1
[input] priority:2 value:2
[input] priority:0 value:3
[input] priority:1 value:4
[input] priority:2 value:5
[input] priority:0 value:6
[input] priority:1 value:7
[input] priority:2 value:8
[input] priority:0 value:9
[input] priority:1 value:10
[input] priority:2 value:11
[input] priority:0 value:12
[input] priority:1 value:13
[input] priority:2 value:14
[output] priority:0 value:0
[output] priority:0 value:6
[output] priority:0 value:12
[output] priority:0 value:3
[output] priority:0 value:9
[output] priority:1 value:13
[output] priority:1 value:10
[output] priority:1 value:4
[output] priority:1 value:1
[output] priority:1 value:7
[output] priority:2 value:2
[output] priority:2 value:5
[output] priority:2 value:14
[output] priority:2 value:11
[output] priority:2 value:8

優先度が同じアイテムは追加した順番で取れていないことが分かります．
当然優先度が同じであれば，どこから取れても良いので動作としては正しいです．
ですが，優先度が同じ場合に追加した順番でアイテム取得したいケースがあったので，実装しました．

アプローチ

ヒープキューを自前で実装するのは手間なので，簡単な方法を取ります．
PythonでのPriorityQueueではキューに追加したアイテム同士を数値として比較できる場合，その数値が低い方が優先度が高くなります．
（数値比較できない場合は先述のようにdataclassesを使います・・・参照）
この性質を利用して，intの優先度の他にキューに追加した時間もキューアイテムに保持して比較対象とします．
そうすることで，キューに追加した順番でアイテムを取得できます．

ついでに最大優先度の取得や，プッシュ・ポップの簡略化も実装します．

実装

以下のようにキューアイテムとキューのクラスを定義します．

import time
from typing import Any, Tuple, Optional
from dataclasses import dataclass, field, asdict
import queue


@dataclass(order=True)
class QueueItem:
    """
    キューアイテム
    """
    priority: int
    item: Any = field(compare=False)
    occurred: float = field(default=0.0, compare=True)


class PriorityQueue(queue.PriorityQueue):
    """
    時間を加味した優先度付きキュークラス
    """

    def __init__(self, priority_name: str = "priority"):
        """
        Args:
            priority_name: 優先度の変数名
        """
        self.priority_name: str = priority_name
        super().__init__()

    @property
    def max_priority(self, priority_name: Optional[str] = None) -> int:
        """
        最大のpriorityを取得
        """
        if not self.queue:
            return 0
        priority_name = priority_name or self.priority_name
        assert hasattr(self.queue[0], priority_name)

        max_priority = max([queue_item.priority for queue_item in self.queue])
        return max_priority

    def put_item(self, **kwargs):
        """
        put override
        """
        # プライオリティ＋登録順で優先順位を決めたい
        if kwargs.get("occurred", None) is None:
            kwargs["occurred"] = time.perf_counter()

        # アイテム作成・追加
        queue_item = QueueItem(**kwargs)
        super().put(queue_item)

    def get_item(self) -> Tuple[Any, ...]:
        """
        get override
        """
        queue_item: QueueItem = super().get()
        item = tuple(asdict(queue_item).values())
        return item

使い方としては以下のような感じです．

def main():
    # キューを用意
    priority_queue = PriorityQueue()
    for index in range(15):
        priority = index % 3
        priority_queue.put_item(priority=priority, item=index)
        print(f"[input] priority:{priority} value:{index}")

    # 取り出し
    while not priority_queue.empty():
        priority, value, *_ = priority_queue.get_item()
        print(f"[output] priority:{priority} value:{value}")

if __name__ == "__main__":
    main()

[input] priority:0 value:0
[input] priority:1 value:1
[input] priority:2 value:2
[input] priority:0 value:3
[input] priority:1 value:4
[input] priority:2 value:5
[input] priority:0 value:6
[input] priority:1 value:7
[input] priority:2 value:8
[input] priority:0 value:9
[input] priority:1 value:10
[input] priority:2 value:11
[input] priority:0 value:12
[input] priority:1 value:13
[input] priority:2 value:14
[output] priority:0 value:0
[output] priority:0 value:3
[output] priority:0 value:6
[output] priority:0 value:9
[output] priority:0 value:12
[output] priority:1 value:1
[output] priority:1 value:4
[output] priority:1 value:7
[output] priority:1 value:10
[output] priority:1 value:13
[output] priority:2 value:2
[output] priority:2 value:5
[output] priority:2 value:8
[output] priority:2 value:11
[output] priority:2 value:14

優先度が同じアイテムはキューに追加した順番で取れているのが分かります．

まとめ

優先度が同じ場合，キューに追加した順番で取得できるキューの実装を紹介しました．
Pythonのオブジェクトの比較方法だけの話なので大げさな話でもないのですが，
単純に優先度という数値に捉われない方法ということで備忘録的に書かせて頂きました．
ご参考になれば幸いです．

余談）
オーバーライドしたプッシュ・ポップの実装は要・不要分かれると思います．
QueueItemで管理しているデータだけ（先述の実装だとitem）取得したい場合は，使える実装かと思います．