mercredi 22 juillet 2020

How to implement ordered class in Python?

I am trying to implement a priority queue class using the heapq library, which requires that items inserted into the queue are comparable with each other.

I would like for the constructor of the queue class to take in a comparator function which, when passed two items, will return a number less than, equal to, or greater than 0 to indicate the relationship between the priorities of the two items. Using this comparator function, I somehow need to create objects to be inserted into the queue which can be directly compared with each other. Here is what I have right now:

from functools import total_ordering
from heapq import heappush, heappop

class PriorityQueue:
    __REMOVED = "<removed-task>" 

    def __init__(self, comparator):
        self._pq = []  # list of entries arranged in a heap
        self._entry_finder = {}  # mapping of tasks to entries

        @total_ordering
        class PrioritizedItem:
            def __init__(self, task):
                self._task = task

            @property
            def task(self):
                return self._task

            @task.setter
            def task(self, task):
                self._task = task

            def __eq__(self, other):
                return comparator(self._task, other._task) == 0

            def __lt__(self, other):
                return comparator(self._task, other._task) < 0

        self._item_class = PrioritizedItem

    def clear(self):
        "Removes all of the elements from this priority queue. The queue will be empty after this call returns."
        self._pq.clear()
        self._entry_finder.clear()

    def remove(self, task):
        "Mark an existing task as REMOVED. Raise KeyError if not found."
        entry = self._entry_finder.pop(task)
        entry.task = self.__REMOVED

    def add(self, task):
        "Add a new task or update the priority of an existing task"
        if task in self._entry_finder:
            self.remove(task)

        entry = this._item_class(task)

        self._entry_finder[task] = entry
        heappush(self._pq, entry)

    def pop(self):
        "Remove and return the lowest priority task. Raise KeyError if empty."
        while self._pq:
            task = heappop(self._pq).task
            if task is not self.__REMOVED:
                del self._entry_finder[task]
                return task
        raise KeyError("pop from an empty priority queue")

I didn't want to make the PrioritizedItem class a standalone class outside of PriorityQueue since then I would need to pass the comparator function into the constructor for each instance of the class. It doesn't really make sense to me to define an ordered class if every instance of the class could potentially contain a different comparator. (For instance, if two instances x and y were created with different comparator functions, then we could have x == y but y != x).

On the other hand, I'm not sure if my current approach is great either since I'm defining a class inside the constructor of the queue and assigning the class to an instance variable, which just feels weird.

Is there any way to improve this code? Thanks!

Aucun commentaire:

Enregistrer un commentaire