Pythonic way of determining if the current element is the first or last element of a generator?

Question

I am going through a generator, whats the Pythonic way of determining if the current element is the first or last element of a generator, given that they need special care?

thanks

basically generating tags, so i have items like

<div class="first">1</div>
<div>...</div>
<div class="last">n</div>

so i would like to keep the last item in loop?

Answer 1

Here's an enumerate-like generator that skips ahead one; it returns -1 for the last element.

>>> def annotate(gen):
...     prev_i, prev_val = 0, gen.next()
...     for i, val in enumerate(gen, start=1):
...         yield prev_i, prev_val
...         prev_i, prev_val = i, val
...     yield '-1', prev_val
>>> for i, val in annotate(iter(range(4))):
...     print i, val
... 
0 0
1 1
2 2
-1 3

It can't tell whether the generator passed to it is "fresh" or not, but it still tells you when the end is nigh:

>>> used_iter = iter(range(5))
>>> used_iter.next()
0
>>> for i, val in annotate(used_iter):
...     print i, val
... 
0 1
1 2
2 3
-1 4

Once an iterator is used up, it raises StopIteration as usual.

>>> annotate(used_iter).next()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in annotate
StopIteration

Answer 2

The way that I do this is similar to some of the other answers here - but I do it this way, as a matter of preference. Maybe it will suit your preferences too.

With my function below, I can write code like this:

values = [10, 11, 12, 13]
for i, val, isfirst, islast in enumerate2(values):
  if isfirst:
    print 'BEGIN...', val
  elif islast:
    print val, '... END'
  else:
    print val

Here is the function definition:

def enumerate2(iterable_):
  it = iter(iterable_)
  try:
    e = it.next()
    isfirst = True
    i = 0
    try:
      while True:
        next_e = it.next()
        yield (i, e, isfirst, False)
        i += 1
        isfirst = False
        e = next_e
    except StopIteration:
      yield (i, e, isfirst, True)
  except StopIteration:
    pass

Answer 3

For the first, use a flag to tell whether or not you've processed any. For the last, hold the next value in a variable, and if there are no more then that's the last one.

Answer 4

Well, as for the first element:

for n, item in enumerate(generator()):
  if n == 0:
    # item is first
# out of the loop now: item is last

Answer 5

Turn it into a sequence, example:

>>> gen = (x for x in range(5))
>>> L = list(gen)
>>> L[0]
0
>>> L[-1]
4
>>>

If you need to do this during the loop:

>>> gen = (x for x in range(5))
>>> L = list(gen)
>>> for idx, item in enumerate(L):
...    if idx == 0:
...        print(u'{item} is first'.format(item=item))
...    if idx == len(L) - 1:
...        print(u'{item} is last'.format(item=item))
...
0 is first
4 is last
>>>

Clearly, this is not the solution, if you are the one who created the generator, and need it to stay that way (for memory savings), but if you don't care, this is more Pythonic per se than setting flags (which is implicit at best, sine it relies on the last element during iteration persisting), and enumerate won't get you any closer to finding the last element.

Answer 6

Of course, it violates all generator's advantages, but if your iterable is not large, you should use:

list(gener)[1:-1]

Answer 7

If you're concerned about potentially large collections built dynamically, so that you don't want to temporarily place it into a single data structure, here's a different way:

FLAGMASK_FIRST = 1
FLAGMASK_LAST = 2

def flag_lastfirst(collection):
    first_flag = FLAGMASK_FIRST
    first = True
    index = 0
    for element in collection:
        if not first:
            yield index, first_flag, current
            index += 1
            first_flag = 0
        current = element
        first = False
    if not first:
        yield index, first_flag | FLAGMASK_LAST, current

l = [1, 2, 3, 4]
for k in flag_lastfirst(l):
    print(k)

The function will produce a sequence of tuples, one for each element from the original collection.

The contents of the tuple:

• t[0] = 0-based index
• t[1] = bitwise flags, FLAGMASK_FIRST is present if the element is the first element, FLAGMASK_LAST is present if the element is the last element
• t[2] = The original element from the original collection

Sample output from the code above:

 +-- 0-based index
 v
(0, 1, 1)
(1, 0, 2)
(2, 0, 3)
(3, 2, 4)
    ^  ^
    |  +-- the element from the original collection
    |
    +-- 1 means first, 2 means last,
        3 means both first and last, 0 is everything else

I'm sure there are nicer ways to build this kinda thing, but this is my contribution anyway.