Python list comprehension with multiple elements of variable length?

Question

I'm writing code to convert a CSV into XML. Suppose I have a single list like:

input = ['name', 'val', 0, \
         'name', 'val', 1, 'tag', 'val', \
         'name', 'val', 2, 'tag', 'val', 'tag', 'val', \
         'name', 'val', 0]

Every slice of this list starting with 'name' denotes an element with a name, a value, and a variable number of optional tag-value pairs.

I want to turn this into:

output = [['name', 'val', []], 
          ['name', 'val', ['tag', 'val']],
          ['name', 'val', ['tag', 'val', 'tag', 'val']],
          ['name', 'val', []]]

No need to separate the tag-value pairs into tuples, this is handled in a separate method. I have a solution, but it's not very pythonic:

output=[]
cursor=0

while cursor < len(input):
    name=input[cursor]
    val=input[cursor+1]
    ntags=int(input[cursor+2])
    optslice=input[cursor+3:cursor+3+2*ntags]
    cursor = cursor+3+2*ntags
    print name, val, ntags, optslice, cursor
    output.append([name, val, optslice])    
print output

> name val 0 [] 3
> name val 1 ['tag', 'val'] 8
> name val 2 ['tag', 'val', 'tag', 'val'] 15
> name val 0 [] 18

> [['name', 'val', []], ['name', 'val', ['tag', 'val']], ['name', 'val', ['tag', 'val', 'tag', 'val']], ['name', 'val', []]]

I think I could probably do this as a list comprehension, but the variable length of each element is throwing me for a loop. The input is parsed from a CSV, and I can change the format to better fit a different solution. Ideas?

Answer 1

I'd use an iterator instead of your cursor and then drive the comprehension with for name in it .

it = iter(input)
output = [[name, next(it), [next(it) for _ in range(2 * next(it))]] for name in it]

Or with islice :

from itertools import islice

it = iter(input)
output = [[name, next(it), list(islice(it, 2 * next(it)))] for name in it]

That said, I suspect you shouldn't have all data in that flat list in the first place. Likely your CSV file has structure that you should use instead. Ie, don't flatten two-dimensional data just so you need to unflatten it back. But your question is interesting nonetheless :-)

Answer 2

I don't know how pythonic you consider this, but you could do something like this

finallist = []
therest = x

while therest:
    name, val, count, *therest = therest
    sublist, therest = rest[:2*count], rest[2*count:]
    finallist.append([name, val] + [sublist])

Answer 3

Here is my code:

data = ['name', 'val', 0,
        'name', 'val', 1, 'tag', 'val',
        'name', 'val', 2, 'tag', 'val', 'tag', 'val',
        'name', 'val', 0]

tmp = [
    [
        data[pos:pos + 2],
        [i for i in data[pos + 3:pos + 3 + data[pos + 2] * 2]]
    ] for pos, e in enumerate(data) if e == 'name']

for e in tmp:
    print e

The output is:

# [['name', 'val'], []]
# [['name', 'val'], ['tag', 'val']]
# [['name', 'val'], ['tag', 'val', 'tag', 'val']]
# [['name', 'val'], []]

Answer 4

If you really want to use pure list comprehension:

a = ['name', 'val', 0, \
              'name', 'val', 1, 'tag', 'val', \
              'name', 'val', 2, 'tag', 'val', 'tag', 'val', \
              'name', 'val', 0]


print(
[grouped[:2] + [tag for tag in grouped[3:]] for grouped in
    [
        a[i:i+(a[i+1:].index("name") + 1 if a[i+1:].count("name") else len(a[i:])+1)]
        for i, x in enumerate(a) if x == "name"
    ]
])

It is really ugly though.