Skip to content

Instantly share code, notes, and snippets.

@nathan-cruz77
Created December 8, 2023 06:23
Show Gist options
  • Save nathan-cruz77/e1535b7fbf6865214ec21aa5aa62a4b3 to your computer and use it in GitHub Desktop.
Save nathan-cruz77/e1535b7fbf6865214ec21aa5aa62a4b3 to your computer and use it in GitHub Desktop.
Enumerate iterables recursively
from collections.abc import Iterable
# Enumerate iterable containing iterables recursively.
#
# Like `enumerate` but can recursively go into nested iterables yielding each
# entry's index as a tuple.
#
#
# Sample usage:
#
# >>> list(enumerate_n('123'))
# [(0, '1'), (1, '2'), (2, '3')]
#
# >>> list(enumerate_n(['123']))
# [(0, '123')]
#
# >>> list(enumerate_n(['123'], n=2))
# [((0, 0), '1'), ((0, 1), '2'), ((0, 2), '3')]
#
# >>> list(enumerate_n(['123', [1, 2]], n=2))
# [((0, 0), '1'),
# ((0, 1), '2'),
# ((0, 2), '3'),
# ((1, 0), 1),
# ((1, 1), 2)]
#
def enumerate_n(iterable, start=0, n=1):
count = start
for item in iterable:
if isinstance(item, Iterable) and n > 1:
for index, value in enumerate_n(iter(item), start=start, n=n - 1):
if not isinstance(index, Iterable):
index = [index]
yield tuple([count, *index]), value
else:
yield count, item
count += 1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment