"""Reference solution for D1 — Resumable Iterator with Serializable State. DO NOT OPEN BEFORE YOU FINISH THE DIAGNOSTIC. -------------------------------------------------------------------------- The idea that makes gates 3 and 4 fall out -------------------------------------------------------------------------- The naive checkpoint is "how many outputs have I emitted." That works for gate 1 and gate 2 and is *wrong* for gate 3 onward, because with a filter in the pipeline the mapping from output count to source position is not invertible without replaying the pipeline anyway. The checkpoint that actually works is a position in the SOURCE plus a position WITHIN the current source item's expansion: (source_pos, sub_index) source_pos -- index of the source item currently being emitted from sub_index -- how many outputs of THAT source item have been emitted Then every transform is the same thing: a function from one source item to a list of zero or more outputs. map(fn) -> item -> [fn(item)] exactly 1 filter(pred) -> item -> [item] or [] 0 or 1 flat_map(fn) -> item -> list(fn(item)) 0 or N Resume is then uniform and needs no special cases: 1. skip `source_pos` items from a fresh source 2. expand the next item through the pipeline 3. drop the first `sub_index` outputs 4. continue normally Gate 4 is free once you have this. That is the point of the problem: the gates are designed so that the naive model passes 1-2 and collapses at 3, and the right model passes all four with no rework. In the real format you do not get to see gate 4 first -- so the skill being measured is choosing a representation that survives requirements you have not been told yet. State size is O(1): two integers and a version tag, regardless of how many items have been consumed. """ from __future__ import annotations from collections import deque from typing import Any, Callable, Iterable, Iterator, List, Tuple SourceFactory = Callable[[], Iterable[Any]] _MISSING = object() class _Pipeline: """An ordered chain of 1 -> N transforms over source items.""" __slots__ = ("_stages",) def __init__(self, stages: Tuple[Tuple[str, Callable], ...] = ()) -> None: self._stages = tuple(stages) def extended(self, kind: str, fn: Callable) -> "_Pipeline": return _Pipeline(self._stages + ((kind, fn),)) def expand(self, item: Any) -> List[Any]: """All outputs produced by a single source item.""" current: List[Any] = [item] for kind, fn in self._stages: nxt: List[Any] = [] if kind == "map": for value in current: nxt.append(fn(value)) elif kind == "filter": for value in current: if fn(value): nxt.append(value) else: # flat_map for value in current: nxt.extend(fn(value)) current = nxt if not current: break return current class ResumableIterator: """An iterator over a replayable source that can checkpoint and resume.""" STATE_VERSION = 1 __slots__ = ( "_source_factory", "_pipeline", "_source_pos", "_sub_index", "_source", "_buffer", "_started", ) def __init__( self, source_factory: SourceFactory, _pipeline: _Pipeline | None = None, _source_pos: int = 0, _sub_index: int = 0, ) -> None: if not callable(source_factory): raise TypeError("source_factory must be a zero-argument callable") self._source_factory = source_factory self._pipeline = _pipeline if _pipeline is not None else _Pipeline() self._source_pos = _source_pos self._sub_index = _sub_index self._source: Iterator[Any] | None = None self._buffer: deque = deque() self._started = False # --- transform chaining ------------------------------------------------ # # Chaining returns a NEW iterator carrying the same position. Transforms # are part of the iterator's identity, not part of its state -- which is # why resume() takes only (factory, state) and you re-declare the # pipeline on the resumed object. def _with(self, kind: str, fn: Callable) -> "ResumableIterator": if self._started: raise RuntimeError("cannot add transforms after iteration has begun") return ResumableIterator( self._source_factory, self._pipeline.extended(kind, fn), self._source_pos, self._sub_index, ) def map(self, fn: Callable[[Any], Any]) -> "ResumableIterator": return self._with("map", fn) def filter(self, pred: Callable[[Any], bool]) -> "ResumableIterator": return self._with("filter", pred) def flat_map(self, fn: Callable[[Any], Iterable[Any]]) -> "ResumableIterator": return self._with("flat_map", fn) # --- iterator protocol ------------------------------------------------- def __iter__(self) -> "ResumableIterator": return self def _start(self) -> None: if self._started: return self._source = iter(self._source_factory()) # Skip source items already fully consumed. next(..., sentinel) so an # exhausted source leaves us cleanly exhausted rather than raising # out of the middle of setup. for _ in range(self._source_pos): if next(self._source, _MISSING) is _MISSING: break self._started = True def __next__(self) -> Any: self._start() assert self._source is not None while not self._buffer: item = next(self._source, _MISSING) if item is _MISSING: raise StopIteration outputs = self._pipeline.expand(item) if self._sub_index: # Only ever non-zero for the first item after a resume. outputs = outputs[self._sub_index :] if outputs: self._buffer.extend(outputs) else: # This source item yielded nothing new -- advance past it. self._source_pos += 1 self._sub_index = 0 value = self._buffer.popleft() self._sub_index += 1 if not self._buffer: self._source_pos += 1 self._sub_index = 0 return value # --- checkpointing ----------------------------------------------------- def state(self) -> dict: return { "version": self.STATE_VERSION, "source_pos": self._source_pos, "sub_index": self._sub_index, } @classmethod def resume(cls, source_factory: SourceFactory, state: dict) -> "ResumableIterator": version = state.get("version") if version != cls.STATE_VERSION: raise ValueError( f"unsupported checkpoint version {version!r}; " f"this build reads version {cls.STATE_VERSION}" ) source_pos = state["source_pos"] sub_index = state["sub_index"] if not isinstance(source_pos, int) or not isinstance(sub_index, int): raise ValueError("checkpoint positions must be integers") if source_pos < 0 or sub_index < 0: raise ValueError("checkpoint positions must be non-negative") return cls(source_factory, None, source_pos, sub_index) # --- batching ---------------------------------------------------------- def batch(self, n: int) -> Iterator[list]: if not isinstance(n, int) or isinstance(n, bool): raise ValueError("batch size must be an int") if n <= 0: raise ValueError("batch size must be positive") return self._batched(n) def _batched(self, n: int) -> Iterator[list]: # A generator, so it stays lazy: nothing is pulled from the source # until the consumer asks for a batch. chunk: list = [] for item in self: chunk.append(item) if len(chunk) == n: yield chunk chunk = [] if chunk: yield chunk