stl_treemap.iterators

  1"""STL-like forward and reverse iterators for tree containers."""
  2
  3from __future__ import annotations
  4
  5from abc import ABC, abstractmethod
  6from typing import Any, Protocol
  7
  8from stl_treemap.tree_node import TreeNode
  9
 10
 11class IterableContainer(Protocol):
 12    """Protocol for containers that support forward and backward node traversal."""
 13
 14    def next(self, node: Any) -> Any:
 15        """Return the node following *node* in forward order."""
 16        ...
 17
 18    def prev(self, node: Any) -> Any:
 19        """Return the node preceding *node* in forward order."""
 20        ...
 21
 22
 23class BaseIterator[K, V, C: IterableContainer](ABC):
 24    """
 25    Base class for STL-like iterators.
 26
 27    References a node and a container. Navigation is achieved by calling
 28    the container's prev() and next() methods.
 29    """
 30
 31    def __init__(self, node: TreeNode[K, V], container: C) -> None:
 32        """
 33        Create an iterator from a node and a container.
 34
 35        Args:
 36            node: Start node.
 37            container: Container to traverse.
 38
 39        """
 40        self._n = node
 41        self._c = container
 42
 43    @abstractmethod
 44    def next(self) -> None:
 45        """Advance to the next node."""
 46
 47    @abstractmethod
 48    def prev(self) -> None:
 49        """Retreat to the previous node."""
 50
 51    def equals(self, rhs: object) -> bool:
 52        """
 53        Check equality of two iterators.
 54
 55        Two iterators are equal when they point to the same node of the same container.
 56
 57        Args:
 58            rhs: Right-hand side iterator.
 59
 60        Returns:
 61            True when both iterators point to the same node of the same container.
 62
 63        Raises:
 64            ValueError: When iterators are of different types.
 65            ValueError: When iterators belong to different containers.
 66
 67        """
 68        if not isinstance(rhs, type(self)):
 69            lhs_class = type(self).__name__
 70            rhs_class = type(rhs).__name__
 71            raise ValueError(f"Can't compare an instance of {lhs_class} with an instance of {rhs_class}")
 72        if self._c is not rhs.container:
 73            raise ValueError("Iterators belong to different containers")
 74        return self._n is rhs.node
 75
 76    @property
 77    def node(self) -> TreeNode[K, V]:
 78        """@private Current node."""
 79        return self._n
 80
 81    @property
 82    def key(self) -> K:
 83        """Key of the current node."""
 84        return self._n.key
 85
 86    @property
 87    def value(self) -> V:
 88        """Value of the current node."""
 89        return self._n.value
 90
 91    @property
 92    def container(self) -> C:
 93        """@private Container that holds the current node."""
 94        return self._c
 95
 96
 97class TreeIterator[K, V, C: IterableContainer](BaseIterator[K, V, C]):
 98    """
 99    STL-like forward iterator.
100
101    More verbose than Python iterators, but allows iteration over any range
102    within the container.
103
104    Example::
105
106        m = TreeMap()
107        it = m.begin()
108        while not it.equals(m.end()):
109            print(f"key: {it.key}, value: {it.value}")
110            it.next()
111    """
112
113    def __init__(self, *args: Any) -> None:
114        """
115        Create a forward iterator.
116
117        Three construction forms are supported:
118
119        1. From a node and container: ``TreeIterator(node, container)``
120        2. Copy of a TreeIterator: ``TreeIterator(other)``
121        3. Convert from ReverseIterator: ``TreeIterator(reverse_iter)``
122
123        Args:
124            *args: Either ``(node, container)`` or a single iterator to copy/convert.
125
126        Raises:
127            ValueError: When arguments are invalid or the source object is unsupported.
128
129        """
130        if len(args) == 2:  # noqa: PLR2004
131            node, container = args
132            super().__init__(node, container)
133        elif len(args) == 1:
134            obj = args[0]
135            if isinstance(obj, TreeIterator):
136                super().__init__(obj._n, obj._c)
137            elif isinstance(obj, ReverseIterator):
138                c = obj._c
139                super().__init__(c.next(obj._n), c)
140            else:
141                raise ValueError(f"Can't create an Iterator from {type(obj).__name__}")
142        else:
143            raise ValueError("Can't create an Iterator with provided parameters")
144
145    def next(self) -> None:
146        """Advance to the next node in the container."""
147        self._n = self._c.next(self._n)
148
149    def prev(self) -> None:
150        """Retreat to the previous node in the container."""
151        self._n = self._c.prev(self._n)
152
153
154class ReverseIterator[K, V, C: IterableContainer](BaseIterator[K, V, C]):
155    """
156    STL-like backward iterator.
157
158    Traverses the container in reverse order: next() moves backward,
159    prev() moves forward.
160
161    Example::
162
163        m = TreeMap()
164        it = m.rbegin()
165        while not it.equals(m.rend()):
166            print(f"key: {it.key}, value: {it.value}")
167            it.next()
168    """
169
170    def __init__(self, *args: Any) -> None:
171        """
172        Create a reverse iterator.
173
174        Three construction forms are supported:
175
176        1. From a node and container: ``ReverseIterator(node, container)``
177        2. Copy of a ReverseIterator: ``ReverseIterator(other)``
178        3. Convert from TreeIterator: ``ReverseIterator(forward_iter)``
179
180        Args:
181            *args: Either ``(node, container)`` or a single iterator to copy/convert.
182
183        Raises:
184            ValueError: When arguments are invalid or the source object is unsupported.
185
186        """
187        if len(args) == 2:  # noqa: PLR2004
188            node, container = args
189            super().__init__(node, container)
190        elif len(args) == 1:
191            obj = args[0]
192            if isinstance(obj, ReverseIterator):
193                super().__init__(obj._n, obj._c)
194            elif isinstance(obj, TreeIterator):
195                c = obj._c
196                super().__init__(c.prev(obj._n), c)
197            else:
198                raise ValueError(f"Can't create an ReverseIterator from {type(obj).__name__}")
199        else:
200            raise ValueError("Can't create a Reverse Iterator with provided parameters")
201
202    def next(self) -> None:
203        """Advance to the previous node in the container (reverse order)."""
204        self._n = self._c.prev(self._n)
205
206    def prev(self) -> None:
207        """Retreat to the next node in the container (reverse order)."""
208        self._n = self._c.next(self._n)