5
\$\begingroup\$

I am new to Python and am trying to see if there is a cleaner way to write my code. The following code take iterables as parameter it produces the first value from the first parameter, then the first value from the the second parameter, ..., then the first value from the last parameter; then the second value from the first parameter and so on.If any iteratable produces no more values, this generator produces no more values.

If I do:

for i in skip('abcde','fg','hijk'):
  print(i,end='')

produces the values 'a', 'f', 'h', 'b', 'g', 'i', and 'c'. The code I've written is as under(which works):

def skip(*args):
  itrs = [iter(arg) for arg in args]
  while itrs != []: 
     temp = [] 
     for i in range(len(itrs)): 
        try: 
          yield next(itrs[i]) 
          temp.append(itrs[i]) 
        except StopIteration: 
          pass
      itrs = temp

Is there a cleaner/ concise way to write the code(without using any external libraries)?

\$\endgroup\$
  • \$\begingroup\$ The code given does not stop at c, it produces afhbgicjdke. I think you meant to return instead of pass in the except block \$\endgroup\$ – XrXrXr Jan 25 '15 at 17:04
  • \$\begingroup\$ The function, as written (but not as specified), is similar to the built-in itertools.zip_longest(), but with no fill values. \$\endgroup\$ – 200_success Jan 25 '15 at 18:52
6
\$\begingroup\$

First off, I think there is an error in the original code. The code produces afhbgicjdke instead of the expected afhbgic. I think the intention was to return instead of pass in the stop iteration block, the answer below work under this assumption.

Let's look at the for loop in the middle

for i in range(len(itrs)):
    try:
        yield next(itrs[i])
        temp.append(itrs[i])
    except StopIteration:
        return

No where in the body operate on the index, so every itrs[i] can be replaced with just i with a simpler for i in iters. This gives us

for i in itrs:
    try:
        yield next(i)
        temp.append(i)
    except StopIteration:
        return

Now lets look at the while body

while itrs != []:
    temp = []
    # ... the for loop that includes an append to itrs
    itrs = temp

This basically gives the code a way to exit when everything is finished. However, if you think about how the code works, the only time itrs has a lower length than args is when one of the itrs reaches StopIteration. In that case, we would have returned already!

So the while loop can become a while True loop and temp and it's usage can be ommited.

def skip(*args):
    itrs = [iter(arg) for arg in args]
    while True:
        for i in itrs:
            try:
                yield next(i)
            except StopIteration:
                return

Now we need to convince ourselves that the code will actually return at some point. This isn't too hard, as long as one of the iterable passed in is finite i.e. has an end, we will reach StopIteration at some point. And in the case that every iterable is infinite, the original code would run forever too. So we are good there.

But there is another problem. What if skip is called with no argument? This will result in us never returning from the function, since there is no iterable for us to trigger StopIteration! In the original code, this is dealt with implicitly through temp being empty. Since we don't have that anymore, we have to handle it explicitly. This gives us the final code:

def skip(*args):
    if len(args) == 0:
        return
    itrs = [iter(arg) for arg in args]
    while True:
        for i in itrs:
            try:
                yield next(i)
            except StopIteration:
                return

With this

for i in skip('abcde','fg','hijk'):
    print(i,end='')

outputs afhbgic and

for i in skip():
    print(i,end='')

is a noop

As an aside, the indentation in the original code is inconsistent, which makes it un-runnable. Please be sure to post usable code next time

\$\endgroup\$
5
\$\begingroup\$

You said about not using external libraries, but if it's OK for you to use the Python standard library, you can pretty much do what you want with cycle function from itertools used on an iterable of iterators:

from itertools import cycle

def skip(*args):
    return (next(i) for i in cycle(iter(arg) for arg in args))

or a more condensed version for Python 3:

from itertools import cycle

def skip(*args):
    return map(next, cycle(map(iter, args)))

The same more condensed version for Python 2:

from itertools import cycle, imap

def skip(*args):
    return imap(next, cycle(imap(iter, args)))

In any case you don't need to catch StopIteration in Python 2 inside generator: throwing it ends the generator automatically and is caught by the caller. So even without itertools you can do something like this:

def skip(*args):
    if not args:
        return
    iters = [iter(arg) for arg in args]
    while True:
        for it in iters:
            yield next(it)

As per comment by @XrXrXr the solution with cycle due to deprecation of bubbling of StopIteration in Python 3.5 would become:

def skip(*args):
    try:
        yield from (next(i) for i in cycle(iter(arg) for arg in args))
    except StopIteration:
        return
\$\endgroup\$
  • 2
    \$\begingroup\$ Note that this behavior of StopIteration is going to be deprecated starting 3.5. See PEP479 \$\endgroup\$ – XrXrXr Jan 25 '15 at 22:29

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.