CODEWITHSUNDEEP
python
Jamie Twells
Jamie Twells

Reputation: 2144

How to do something to an iterator, then something else to the next iterator in python

What's the 'correct' python way of doing something to the first element, but something else to the second, and something different again for the third in a list, then repeat, eg:

a = [2, "foo1", "bar1", 5, "foo3", "bar2", 3, "foo2", "bar3"]
my_function(a)

should give

20
foo1
foobar1
50
foo3
foobar2
30
foo2
foobar3

where "my_function" would be something like:

def my_function(a):
    i = 0
    for line in a:
        if i == 0:
            line = line*10
        if i == 2:
            line = "foo"+line
        i = i + 1
        if i == 3:
            i = 0
        print line

But this looks very 'unpython'. Is there a better way of doing this, without an int to keep track? A way of saying, the first time I call this function, do this but the second time, do this, then the third time, do this, then go back to the beginning and do what you did first. A function that keeps track of how many times it has been called.

Upvotes: 0

Views: 123

Answers (7)

Kei Minagawa
Kei Minagawa

Reputation: 4521

This is C like way rather than Pythonic. So if you know C you can understand it just a few seconds.

def my_function(a):
    for n, line in enumerate(a):
        if n % 3 == 0:
            line = line * 10
        elif n % 3 == 2:
            line = "foo" + line
        print line

Upvotes: 0

Janne Karila
Janne Karila

Reputation: 25207

Working with an iterator, you can do like this:

def my_function(a):
    a = iter(a)
    while True:
        yield 10 * next(a)    
        yield next(a)
        yield "foo" + next(a)

a = [2, "foo1", "bar1", 5, "foo3", "bar2", 3, "foo2", "bar3"]    
print list(my_function(a))
#prints [20, 'foo1', 'foobar1', 50, 'foo3', 'foobar2', 30, 'foo2', 'foobar3']

If anyone is wondering what happens at the end of the list, next(a) will raise StopIteration. The exception terminates the generator my_function, but code that iterates over it -- in this case list() -- will recognize it as a normal end of iteration.

Upvotes: 3

mkriheli
mkriheli

Reputation: 1856

Use itertools for this, I find it most elegant and resilient to the number of elements/actions:

import itertools


def my_function(a):
    actions = itertools.cycle(
        (lambda l: l * 10, lambda l: l, lambda l: "foo" + l)
    )

    for val, action in itertools.izip(a, actions):
        yield action(val)

a = [2, "foo1", "bar1", 5, "foo3", "bar2", 3, "foo2", "bar3", 7]
print list(my_function(a))

Result:

[20, 'foo1', 'foobar1', 50, 'foo3', 'foobar2', 30, 'foo2', 'foobar3', 70]

As you can see from the example, this will work even if the elements count is not a multiplication of the actions count (or even less than actions count).

This is python2 version, if you need 3 please specify.

Upvotes: 0

Aashish P
Aashish P

Reputation: 1966

Assuming a list as an input to the function,

def my_function(a):
   a[::3] = map(lambda x: 10*x, a[::3])
   a[2::3] = map(lambda x: "foo"+x, a[2::3])

Here I am making use of extended slices in python. This is more of a pythonic way i think.

Upvotes: 0

Marcin
Marcin

Reputation: 239000

Other option to the existing ones using compination of zip and iter (explenation here):

def my_function2(a):
    for v1,v2,v3 in zip(*[iter(a)]*3):
        print(v1*10,v2, 'foo'+ v3)

% gives
20 foo1 foobar1
50 foo3 foobar2
30 foo2 foobar3

Upvotes: 0

jonrsharpe
jonrsharpe

Reputation: 122159

You could use something like: 

def my_function (a):
    actions = [lambda l: l * 10,
               lambda l: l,
               lambda l: "foo" + l]
    for i, line in enumerate(a):
        line = actions[i % 3](line)
        print line

Upvotes: 0

mgilson
mgilson

Reputation: 310307

I might do something like:

def my_function(lst):
    items = (lst[i:i+3] for i in xrange(0, len(lst), 3))
    for group in items:
       yield group[0] * 10
       yield group[1]
       yield 'foo' + group[2]

Running on your input:

>>> list(my_function(a))
[20, 'foo1', 'foobar1', 50, 'foo3', 'foobar2', 30, 'foo2', 'foobar3']

I've made 2 big assumptions here -- that your list is indeed a list (or at least a sequence), and that the length of the sequence is divisible by 3 (otherwise you'll end up with an IndexError). Both of these assumptions could be taken care of without too much effort, but I'll leave that as an exercise if you're really interested ;-)

Upvotes: 1

Related Questions