Search code examples
pythonlist-comprehension

python change for_loop in list comprehesion


I've got the code as follows:

ix=0
list_test=[]
for el in parse_data_2['var_nm_final']:    
    if el=='URI':
        ix+=1
        list_test.append(ix)
    else:
        list_test.append(None)

parse_data_2 is my DF. As an output I would like to receive list with incremented value of ix or None depending on my condition. Meaning something like this

1
None
None
None
None
2
None
3
None
None
None
None
4

... etc.

I've tried to convert this loop to list comprehesion like this:

[ix+=1 if el=='URI'else None for el in parse_data_2['var_nm_final']]         

but got error

[ix+=1 if el=='URI'else None for el in parse_data_2['var_nm_final']]
       ^
SyntaxError: invalid syntax

Could you explain me the issue with my code?


Solution

  • Use next on an itertools.count (or just range with generous upper bound):

    >>> parse_data_2 = {'var_nm_final': ["URI", "foo", "bar", "URI", "URI", "blub", "URI"]}
    >>> import itertools
    >>> cnt = itertools.count(1)
    >>> [next(cnt) if el == "URI" else None for el in parse_data_2["var_nm_final"]]
    [1, None, None, 2, 3, None, 4]