Controlling descent in an iterator

hierarchical-iter.py

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Generator, Mapping, Sequence


@dataclass
class Hier[T]:
    things: Mapping[T, Sequence[T]]
    descend = True

    def __iter__(self) -> Generator[T]:
        for cat, children in self.things.items():
            yield cat
            if self.descend:
                yield from children
            self.descend = True


things = {"a": ("a.1", "a.2"), "b": ("b.1", "b.2")}


print("all:")
for thing in Hier(things):
    print(thing)

print()

print("no a.*:")
for thing in (hier := Hier(things)):
    if thing == "a":
        hier.descend = False
    print(thing)