Added __iter__ method

Author: yusinv

CHANGELOG.rst

@@ -28,3 +28,4 @@ Changelog
 
 * Refactoring
 * Added PrioritySearchSet class
+* Added __iter__ method

docs/reference/priority_search_tree_print.rst

@@ -1,5 +1,5 @@
 priority_search_tree.print_helpers
-====================
+==================================
 
 .. testsetup::
 

src/priority_search_tree/ps_set.py

@@ -1,6 +1,7 @@
-import collections
 from typing import Callable
 from typing import Iterable
+from typing import Iterator
+from typing import MutableSet
 from typing import Optional
 from typing import TypeVar
 
@@ -11,13 +12,15 @@
 _V = TypeVar("_V")
 
 
-class PrioritySearchSet(collections.abc.MutableSet):
+class PrioritySearchSet(MutableSet):
     """Mutable Set that maintains priority search tree properties.
 
     PrioritySearchSet can be used to store any type of objects.
     2 functions should be passed to PrioritySearchSet constructor:
-        * ``key_func`` to extract **key** for the object
-        * ``priority_func`` to extract **priority** for the object
+
+    * ``key_func`` to extract **key** for the object
+    * ``priority_func`` to extract **priority** for the object
+
     extracted **key**, **priority** values will be used in underlying PrioritySearchTree
 
     Example::
@@ -138,7 +141,7 @@ def remove(self, value: _V) -> None:
         del self._pst[key]
         del self._values[key]
 
-    def query(self, left: _V, right: _V, bottom: _V) -> list[_V]:
+    def query(self, left: _V, right: _V, bottom: _V) -> list:
         """Performs 3 sided query on PSS.
 
         This function returns list of items that meet the following criteria:
@@ -162,7 +165,7 @@ def query(self, left: _V, right: _V, bottom: _V) -> list[_V]:
         priority_bottom = self.priority_func(bottom)
         return [self._values[x] for x in self._pst.query(key_left, key_right, priority_bottom)]
 
-    def sorted_query(self, left: _V, right: _V, bottom: _V, items_limit: int = 0) -> [_V]:
+    def sorted_query(self, left: _V, right: _V, bottom: _V, items_limit: int = 0) -> list:
         """Performs sorted 3 sided query on PSS.
 
         This function returns list of items that meet the following criteria:
@@ -242,5 +245,11 @@ def discard(self, value) -> None:
             del self._pst[key]
             del self._values[key]
 
-    def __iter__(self):
-        raise NotImplementedError
+    def __iter__(self) -> Iterator:
+        """Create an iterator that iterates values in sorted by **key** order
+
+        Returns:
+            Iterator: in order iterator
+        """
+        for key in self._pst:
+            yield self._values[key]

src/priority_search_tree/ps_tree.py

@@ -1,5 +1,6 @@
-import collections.abc
 from typing import Iterable
+from typing import Iterator
+from typing import MutableMapping
 from typing import Optional
 from typing import Tuple
 from typing import TypeVar
@@ -10,7 +11,7 @@
 _PRIORITY = TypeVar("_PRIORITY")
 
 
-class PrioritySearchTree(collections.abc.MutableMapping):
+class PrioritySearchTree(MutableMapping):
     """Class that represents Priority search tree.
 
     PrioritySearchTree is a mutable mapping that stores **keys** and corresponding **priorities**.
@@ -241,7 +242,7 @@ def _transplant(self, u: Node, v: Node) -> None:
         else:
             u.parent.set_right(v)
 
-    def query(self, key_left: _KEY, key_right: _KEY, priority_bottom: _PRIORITY) -> list[_KEY]:
+    def query(self, key_left: _KEY, key_right: _KEY, priority_bottom: _PRIORITY) -> list:
         """Performs 3 sided query on PST.
 
         This function returns list of items that meet the following criteria:
@@ -283,7 +284,7 @@ def _query_node(node) -> None:
         _query_node(self._root)
         return result
 
-    def sorted_query(self, key_left: _KEY, key_right: _KEY, priority_bottom: _PRIORITY, items_limit: int = 0) -> [_KEY]:
+    def sorted_query(self, key_left: _KEY, key_right: _KEY, priority_bottom: _PRIORITY, items_limit: int = 0) -> list:
         """Performs 3 sided query on PST.
 
         This function returns list of items that meet the following criteria:
@@ -613,5 +614,24 @@ def __getitem__(self, key: _KEY) -> _PRIORITY:
 
         return heap_node.heap_key[0]
 
-    def __iter__(self):
-        raise NotImplementedError
+    def __iter__(self) -> Iterator:
+        """Create an iterator that iterates **keys** in sorted order
+
+        Returns:
+            Iterator: in order iterator
+        """
+        stack = []
+        current = self._root
+        yielded_key = None
+        while True:
+            if current != Node.NULL_NODE:
+                stack.append(current)
+                current = current.left
+            elif stack:
+                current = stack.pop()
+                if current.tree_key != yielded_key:
+                    yielded_key = current.tree_key
+                    yield yielded_key
+                current = current.right
+            else:
+                break

tests/manual/performance.py

@@ -10,7 +10,7 @@
 from priority_search_tree import PrioritySearchSet
 
 NUMBER_OF_ITEMS = [1, 100, 10000, 1000000]
-REPEAT_COUNT = 1
+REPEAT_COUNT = 10
 
 TIMEOUT_SEC = 60
 

tests/test_priority_search_set.py

@@ -12,41 +12,44 @@ def __init__(self, x: int, y: int):
 
     items = [Point(1, 1), Point(2, 2), Point(3, 3), Point(4, 4), Point(5, 6), Point(6, 6)]
 
-    pst = PrioritySearchSet(key_func=lambda v: v.x, priority_func=lambda v: v.y, iterable=items)
-    assert_rb_tree(pst._pst._root)
+    pss = PrioritySearchSet(key_func=lambda v: v.x, priority_func=lambda v: v.y, iterable=items)
+    assert_rb_tree(pss._pst._root)
 
-    assert pst.query(Point(1, 1), Point(2, 2), Point(2, 2)) == [items[1]]
-    assert pst.query(Point(1, 1), Point(5, 1), Point(1, 6)) == [items[4]]
-    assert pst.sorted_query(Point(1, 1), Point(6, 1), Point(1, 6)) == [items[5], items[4]]
-    assert pst.sorted_query(Point(1, 1), Point(4, 1), Point(1, 1), items_limit=1) == [items[3]]
+    assert pss.query(Point(1, 1), Point(2, 2), Point(2, 2)) == [items[1]]
+    assert pss.query(Point(1, 1), Point(5, 1), Point(1, 6)) == [items[4]]
+    assert pss.sorted_query(Point(1, 1), Point(6, 1), Point(1, 6)) == [items[5], items[4]]
+    assert pss.sorted_query(Point(1, 1), Point(4, 1), Point(1, 1), items_limit=1) == [items[3]]
 
-    assert items[2] in pst
-    assert Point(10, 10) not in pst
-    assert Point(2, 10) in pst
+    assert items[2] in pss
+    assert Point(10, 10) not in pss
+    assert Point(2, 10) in pss
 
-    pst.remove(items[2])
-    pst.remove(Point(2, 10))
-    pst.discard(Point(4, 1))
-    pst.discard(Point(4, 1))
+    pss.remove(items[2])
+    pss.remove(Point(2, 10))
+    pss.discard(Point(4, 1))
+    pss.discard(Point(4, 1))
 
     with pytest.raises(KeyError, match="Key not found:"):
-        pst.remove(Point(7, 1))
+        pss.remove(Point(7, 1))
 
     with pytest.raises(KeyError, match="Key not found:"):
-        pst.remove(Point(7, 7))
+        pss.remove(Point(7, 7))
 
-    assert_rb_tree(pst._pst._root)
+    assert_rb_tree(pss._pst._root)
 
-    assert pst.get_with_max_priority().y == 6
+    assert pss.get_with_max_priority().y == 6
 
-    assert pst.pop().y == 6
-    assert pst.pop().y == 6
-    assert pst.pop().y == 1
+    assert pss.pop().y == 6
+    assert pss.pop().y == 6
+    assert pss.pop().y == 1
 
-    assert not pst
+    assert not pss
 
-    pst = PrioritySearchSet(key_func=lambda v: v.x, priority_func=lambda v: v.y)
-    pst.add(Point(1, 1))
-    assert len(pst) == 1
-    pst.clear()
-    assert not pst
+    pss = PrioritySearchSet(key_func=lambda v: v.x, priority_func=lambda v: v.y)
+    pss.add(Point(1, 4))
+    assert len(pss) == 1
+    for p in pss:
+        assert p.x == 1
+        assert p.y == 4
+    pss.clear()
+    assert not pss

tests/test_priority_search_tree.py

@@ -37,6 +37,8 @@ def test_empty_pst():
     assert len(pst) == 0
     with pytest.raises(KeyError):
         pst.get_with_max_priority()
+    for x in pst:
+        raise AssertionError(x)
     pst[1] = 5
     assert pst.pop(1) == 5
     with pytest.raises(KeyError, match="Key not found:"):
@@ -113,6 +115,13 @@ def test_query():
                 assert pst.sorted_query(x_min, x_max, y_min) == query_expected
 
 
+def test_iterator():
+    items = [(0, 0), (1, 6), (6, 3), (7, 5), (8, 8), (2, 1), (3, 7), (4, 4), (5, 2)]
+    pst = PrioritySearchTree(items)
+    for e, k in enumerate(pst):
+        assert e == k
+
+
 def test_sorted_query_limit():
     items = [(0, 0), (1, 6), (2, 2), (3, 7), (4, 4), (5, 2), (6, 3), (7, 5), (8, 8)]
     pst = PrioritySearchTree(items)