stl_treemap.py_iterators

  1"""Python-style forward and backward iterators wrapping tree containers."""
  2
  3from __future__ import annotations
  4
  5from collections.abc import Iterator
  6from typing import Any
  7
  8
  9class PyIterator[T](Iterator):
 10    """
 11    Forward iterator for tree containers.
 12
 13    Implements the Python iterator protocol so it can be used directly in
 14    ``for`` loops.  The container must provide ``py_begin()``, ``py_end()``,
 15    ``next(node)``, ``prev(node)``, and a ``value_policy`` attribute.
 16
 17    Example::
 18
 19        m = TreeMap()
 20        for key, value in m:
 21            print(f"key: {key}, value: {value}")
 22
 23        for value in m.values():
 24            print(f"value: {value}")
 25
 26    """
 27
 28    def __init__(self, container: Any, value_policy: Any = None) -> None:
 29        """
 30        Construct a forward iterator for container.
 31
 32        Args:
 33            container: Container to traverse.
 34            value_policy: Policy controlling what each element yields; defaults
 35                to ``container.value_policy``.
 36
 37        """
 38        self.container = container
 39        self.value_policy = container.value_policy if value_policy is None else value_policy
 40        self.node = container.py_begin()
 41
 42    def __iter__(self) -> PyIterator[T]:
 43        """Return self to satisfy the iterator protocol."""
 44        return self
 45
 46    def __next__(self) -> T:
 47        """Return the next value and advance, or raise StopIteration when exhausted."""
 48        if self.node == self.container.py_end():
 49            raise StopIteration
 50        value: T = self.value_policy.fetch(self.node)
 51        self.node = self.container.next(self.node)
 52        return value
 53
 54    def backwards(self) -> PyReverseIterator[T]:
 55        """Return a reverse iterator for the same container and value policy."""
 56        return PyReverseIterator(self.container, self.value_policy)
 57
 58
 59class PyReverseIterator[T](Iterator):
 60    """
 61    Backward iterator for tree containers.
 62
 63    Implements the Python iterator protocol so it can be used directly in
 64    ``for`` loops.  The container must provide ``py_rbegin()``, ``py_rend()``,
 65    ``next(node)``, ``prev(node)``, and a ``value_policy`` attribute.
 66
 67    Example::
 68
 69        m = TreeMap()
 70        for key, value in m.backwards():
 71            print(f"key: {key}, value: {value}")
 72
 73        for key in m.keys().backwards():
 74            print(f"key: {key}")
 75
 76    """
 77
 78    def __init__(self, container: Any, value_policy: Any = None) -> None:
 79        """
 80        Construct a backward iterator for container.
 81
 82        Args:
 83            container: Container to traverse.
 84            value_policy: Policy controlling what each element yields; defaults
 85                to ``container.value_policy``.
 86
 87        """
 88        self.container = container
 89        self.value_policy = container.value_policy if value_policy is None else value_policy
 90        self.node = container.py_rbegin()
 91
 92    def __iter__(self) -> PyReverseIterator[T]:
 93        """Return self to satisfy the iterator protocol."""
 94        return self
 95
 96    def __next__(self) -> T:
 97        """Return the next value (in reverse order) and advance, or raise StopIteration."""
 98        if self.node == self.container.py_rend():
 99            raise StopIteration
100        value: T = self.value_policy.fetch(self.node)
101        self.node = self.container.prev(self.node)
102        return value
103
104    def backwards(self) -> PyIterator[T]:
105        """Return a forward iterator for the same container and value policy."""
106        return PyIterator(self.container, self.value_policy)