Skip to content

Instantly share code, notes, and snippets.

@GreyElaina
Last active May 22, 2021 16:34
Show Gist options
  • Select an option

  • Save GreyElaina/f9a5f998ec1c3fc7bddb811ce046d0ca to your computer and use it in GitHub Desktop.

Select an option

Save GreyElaina/f9a5f998ec1c3fc7bddb811ce046d0ca to your computer and use it in GitHub Desktop.
Nestable Iterable
class NestableIterable(Iterable[T]):
index_stack: list
iterable: Iterable[T]
def __init__(self, iterable: Iterable[T]) -> None:
self.iterable = iterable
self.index_stack = [0]
def __iter__(self):
index = self.index_stack[-1]
self.index_stack.append(self.index_stack[-1])
start_offset = index + int(bool(index))
try:
for self.index_stack[-1], content in enumerate(
itertools.islice(self.iterable, start_offset, None, None),
start=start_offset,
):
yield content
finally:
self.index_stack.pop()
def with_new(self, target):
self.iterable = target
return self
T = TypeVar("T")
I = TypeVar("I")
class NestableIterable(Generic[I, T]):
iterable: Iterable[T]
indexes: List[I]
generator_with_index_factory: Callable[
[Iterable[T], I], Generator[None, None, Tuple[I, T]]
]
index_increase_func: Callable[[I, Iterable[T]], I]
is_index_origin: Callable[[I], bool]
@staticmethod
def default_generator_factory(iterable: Iterable[T], start: I):
return enumerate(iterable[start:], start=start)
def __init__(
self,
iterable: Iterable[T],
generator_with_index_factory: Callable[
[Iterable[T], I], Generator[None, None, Tuple[I, T]]
] = None,
index_increase_func: Callable[[I, Iterable[T]], I] = lambda x, _: x + 1,
initial_index_value_factory: Callable[[], I] = lambda: 0,
is_index_origin: Callable[[I], bool] = lambda x: x == 0,
) -> None:
self.iterable = iterable
self.indexes = [initial_index_value_factory()]
self.generator_with_index_factory = (
generator_with_index_factory or self.default_generator_factory
)
self.index_increase_func = index_increase_func
self.is_index_origin = is_index_origin
def __iter__(self):
current_index = self.indexes[-1]
self.indexes.append(current_index)
if self.is_index_origin(current_index):
start_offset = current_index
else:
start_offset = self.index_increase_func(current_index, self.iterable)
for self.indexes[-1], content in self.generator_with_index_factory(
self.iterable,
start=start_offset,
):
yield content
self.indexes.pop()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment