INIT

diplomacy/utils/priority_dict.py

@@ -0,0 +1,109 @@
+# ==============================================================================
+# Copyright (C) 2019 - Philip Paquette, Steven Bocco
+#
+#  This program is free software: you can redistribute it and/or modify it under
+#  the terms of the GNU Affero General Public License as published by the Free
+#  Software Foundation, either version 3 of the License, or (at your option) any
+#  later version.
+#
+#  This program is distributed in the hope that it will be useful, but WITHOUT
+#  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+#  FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more
+#  details.
+#
+#  You should have received a copy of the GNU Affero General Public License along
+#  with this program.  If not, see <https://www.gnu.org/licenses/>.
+# ==============================================================================
+""" Priority Dict implementation """
+import heapq
+
+# ------------------------------------------------
+# Adapted from (2018/03/14s):
+# https://docs.python.org/3.6/library/heapq.html#priority-queue-implementation-notes
+# Unlicensed
+class PriorityDict(dict):
+    """ Priority Dictionary Implementation """
+
+    def __init__(self, **kwargs):
+        """ Initialize the priority queue.
+
+            :param kwargs: (optional) initial values for priority queue.
+        """
+        self.__heap = []  # Heap for entries. An entry is a triple (priority value, key, valid entry flag (boolean)).
+        # Dict itself maps key to entries. We override some dict methods (see __getitem__() below)
+        # to always return priority value instead of entry as dict value.
+        dict.__init__(self)
+        for key, value in kwargs.items():
+            self[key] = value
+
+    def __setitem__(self, key, val):
+        """ Sets a key with his associated priority
+
+            :param key: The key to set in the dictionary
+            :param val: The priority to associate with the key
+            :return: None
+        """
+        if key in self:
+            del self[key]
+        # Create entry with val, key and a boolean indicating that entry is valid (True).
+        entry = [val, key, True]
+        dict.__setitem__(self, key, entry)
+        heapq.heappush(self.__heap, entry)
+
+    def __delitem__(self, key):
+        """ Removes key from dict and marks associated heap entry as invalid (False).
+            Raises KeyError if not found.
+        """
+        entry = self.pop(key)
+        entry[-1] = False
+
+    def __getitem__(self, key):
+        """ Returns priority value associated to key. Raises KeyError if key not found. """
+        return dict.__getitem__(self, key)[0]
+
+    def __iter__(self):
+        """ Iterator over all keys based on their priority. """
+
+        def iterfn():
+            """ Iterator """
+            copy_of_self = self.copy()
+            while copy_of_self:
+                _, key = copy_of_self.smallest()
+                del copy_of_self[key]
+                yield key
+
+        return iterfn()
+
+    def smallest(self):
+        """ Finds the smallest item in the priority dict
+
+            :return: A tuple of (priority, key) for the item with the smallest priority, or None if dict is empty.
+        """
+        while self.__heap and not self.__heap[0][-1]:
+            heapq.heappop(self.__heap)
+        return self.__heap[0][:2] if self.__heap else None
+
+    def setdefault(self, key, d=None):
+        """ Sets a default for a given key """
+        if key not in self:
+            self[key] = d
+        return self[key]
+
+    def copy(self):
+        """ Return a copy of this priority dict.
+
+            :rtype: PriorityDict
+        """
+        return PriorityDict(**{key: entry[0] for key, entry in dict.items(self)})
+
+    def keys(self):
+        """ Make sure keys() iterates on keys based on their priority. """
+        return self.__iter__()
+
+    def values(self):
+        """ Makes sure values() iterates on priority values (instead of heap entries) from smallest to highest. """
+        return (self[k] for k in self)
+
+    def items(self):
+        """ Makes sure items() values are priority values instead of heap entries. """
+        return ((key, self[key]) for key in self)