Stream generator. Using an iterator in detail

Viewed 215

I'm trying to figure out how an iterator works with this example:

There is the function that produces stream generator for given iterable object (list, generator, etc) whose elements contain position and value and sorted by order of apperance. Stream generator is equal to initial stream (without position), gaps filled with zeroes.

from itertools import count

def gen_stream(total, sorted_iterable, extractor=lambda x: x):
    sorted_iterator = iter(sorted_iterable)
    iterable = count() if total is None else range(total)
    try:
        current_extracted_record = extractor(next(sorted_iterator))
    except StopIteration:
        current_extracted_record = None
    for i in iterable:
        if current_extracted_record:
            if i == current_extracted_record[0]:
                try:
                    yield current_extracted_record[1]
                    current_extracted_record = extractor(next(sorted_iterator))
                except StopIteration:
                    current_extracted_record = None
            else:
                yield 0
        else:
            yield 0

For example:

gen = gen_stream(9,[(4,111),(7,12)])
list(gen) 
[0, 0, 0, 0, 111, 0, 0, 12, 0] # first element has zero index, so 111 located on fifth position, 12 located on 8th position

This function also support custom position-value extractor for more advanced cases, e.g.

def day_extractor(x):
    months = [31, 28, 31, 30, 31, 31, 30, 31, 30, 31, 30, 31]
    acc = sum(months[:x[1] - 1]) + x[0] - 1
    return acc, x[2]
precipitation_days = [(3,1,4),(5,2,6)]
list(gen_stream(59,precipitation_days,day_extractor)) #59: January and February to limit output
[0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

precipitation_days format is following: (d,m,mm), where d - day in month, m - month, mm - precipitation in millimeters So, in example:

(3,1,4) # January,3 precipitation: 4 mm
(5,2,6) # February,5 precipitation: 6 mm

Extractor passed as optional third parameter with default value - lambda function that handles (position, value) pairs like in first example.

This is where the questions begin

Question 1 Can I replace

    try:
        current_extracted_record = extractor(next(sorted_iterator))
    except StopIteration:
        current_extracted_record = None

with default value for the function next instead of catching the exception StopIteration by using one line of code

current_extracted_record = extractor(next((sorted_iterator), None))

will it always work correctly in other situations?

Question 2 How can this part be replaced by using default value for the method next() and loop while instead of the loop for. In theory, the code should be shorter.

    for i in iterable:
        if current_extracted_record:
            if i == current_extracted_record[0]:
                try:
                    yield current_extracted_record[1]
                    current_extracted_record = extractor(next(sorted_iterator))
                except StopIteration:
                    current_extracted_record = None
            else:
                yield 0
        else:
            yield 0

Question 3 This may seem like an idiotic question, but as I understand it, the extractor does not have an index. Then what do the numbers in square brackets mean?

current_extracted_record[0] 
current_extracted_record[1]

Thanks if you can help.

I apologize for the 3 questions in the thread, but it seems to me that they describe the same problem in different details.

Answer (Question1&Question2)

def gen_stream(total, sorted_iterable, extractor=lambda x: x):
    elem_iter = iter(map(extractor, sorted_iterable))
    pos, val = next(elem_iter, (None, None))
    cnt = 0
    while total is None or cnt < total:
        if cnt == pos:
            yield val
            pos, val = next(elem_iter, (None, None))
        else:
            yield 0
        cnt += 1
3 Answers

Question 1

No, extractor(next(sorted_iterator, None)) will fail if next returns None and extractor can't handle it. For example your day_extractor can't (it crashes) and the default identity extractor doesn't return a pair of index and value, so the later code will fail.

Question 2

You can make the whole thing shorter by mapping extractor on sorted_iterable and asking for next values of that:

def gen_stream(total, sorted_iterable, extractor=lambda x: x):
    specials = map(extractor, sorted_iterable)
    iterable = count() if total is None else range(total)
    current_extracted_record = next(specials, None)
    for i in iterable:
        if current_extracted_record and i == current_extracted_record[0]:
            yield current_extracted_record[1]
            current_extracted_record = next(specials, None)
        else:
            yield 0

Since you don't have negative indexes, we can also use [-1] instead of None to shorten the inner if condition (or with really anything not in count(), e.g., [None]). And let me rename a bit as well.

def gen_stream(total, sorted_iterable, extractor=lambda x: x):
    specials = map(extractor, sorted_iterable)
    indexes = count() if total is None else range(total)
    next_special = next(specials, [-1])
    for i in indexes:
        if i == next_special[0]:
            yield next_special[1]
            next_special = next(specials, [-1])
        else:
            yield 0

Question 3

current_extracted_record is what the extractor returned, i.e., the next special (index, value) pair. Which has indexes 0 and 1.

I would reduce the entire function to

from itertools import count, islice, chain, repeat


def gen_stream(sorted_iterable, *, extractor=lambda x: x, total=None):
    itr = chain(map(extractor, sorted_iterable), repeat(None))

    current = next(itr)
    for i in islice(count(), total):
        if current is None or i != current[0]:
            yield 0
        else:
            yield current[1]
            current = next(itr)

We create an iterator whose elements are the return values of extractor applied to the original iterable, followed by None for as long as necessary. No need to check for StopIteration explicitly.

If total is None, then islice(count(), total) is equivalent to count(); otherwise, it's equivalent to range(total).


Alternatively,

def gen_stream(sorted_iterable, *, extractor=lambda x: x, total=None):
    def stream():
        pos = 0
        for new_pos, value in map(extractor, sorted_iterable):
            yield from repeat(0, new_pos - pos)
            yield value
            pos = new_pos + 1
    yield from islice(chain(stream(), repeat(0)), total)

The internal generator will patch the gaps with 0s, then the final yield from can produce the desired stream.


In each case, I made the two optional arguments keyword-only, since I kept typing gen_stream([...], 9), accidentally setting the extractor instead of the total. You can keep your original signature, or some other variation, if you like.

Your first question is asking if the two-argument form of next can help you simplify the code:

try:
    current_extracted_record = extractor(next(sorted_iterator))
except StopIteration:
    current_extracted_record = None

My answer is a sort of maybe. You can't just put in None as the second argument to the existing next call because there's no guarantee the extractor function can handle being given None as an input. The current code won't ever pass None in to extractor, and while None is supported by the default NOOP function, other extractor functions (like your day index one) may not be happy about it. Now, you could change the specification of the main function to require that any extractor support None (e.g. by passing it through unchanged), but without that you'd need at least a little bit of complexity to handle None. Here's how I'd do it:

raw_record = next(sorted_iterator, None)
current_extracted_record = extractor(raw_record) if raw_record is not None else None

Your second question asks if you can simplify your loop somehow. I don't think the loop itself has any obvious simplifications available (using a while instead of for is probably not productive), but you can simplify the two if statements you have into a single one, since they both do the same thing if the condition is not met (yield 0):

for i in iterable:
    if current_extracted_record and i == current_extracted_record[0]: # combine both conditions
        yield current_extracted_record[1]
        raw_record = next(sorted_iterator, None)      # might as well use the code from q1 here
        current_extracted_record = extractor(raw_record) if raw_record is not None else None
    else:
        yield 0

As for your third question, I'm not sure why you think current_extracted_record won't be indexable. It's a 2-tuple, so there are always going to be two items that can be accessed. In the code, they're an index and a value. With the default extractor, they're directly from the input list (e.g. (4, 111)), but with a custom extractor they can be computed from that original value, rather than being directly a part of it (e.g. the acc and x[2] values produced by your day_extractor function).

Related