Unpack a given number of items in Python?

Question

Is there any way to use the 'splat' operator (e.g. a, *rest = somelist) in such a way that it consumes a certain number of items?

Use case: I want to split some input I'm getting into a number, a list of lists, another number, and another list of lists.

My input looks like this:

5
1 2 3 4
5 6 7 8
9 10 11 12
13 14 15 16
5
1 2 3 4
5 6 7 8
9 10 11 12
13 14 15 16

And I want the names first_num, first_arrangement, second_num, second_arrangement such that:

first_num == 5
first_arrangement == [[1, 2, 3, 4], [5, 6, 7, 8], ...]

and so on.

To do this, it'd be useful to be able to consume a set number of items from the iterable I've got yielding the lines. Something like this would be ideal as an intermediate step: first_num, *[4]first_arrangement, second_num, *[4]second_arrangement = lines

What's the normal/canonical/Pythonic way of solving this?

Answer 1

I think the canonical, pythonic way to do this would be to put the onus on the generator that you're iterating over. I'd define a generator function like so:

import itertools

def generate_arrangements(iterable, size=4):
    iterator = iter(iterable)
    while True:
        yield next(iterator)
        yield list(list(row) for row in itertools.islice(iterator, size))

Say you have your data in a list like so:

data = [
    5,
    [1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12],
    [13, 14, 15, 16],
    5,
    [1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12],
    [13, 14, 15, 16]
]

Then writing:

first_num, first_arr, second_num, second_arr = generate_arrangements(data)

gives you your desired output:

>>> first_num
5
>>> first_arr
[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]

You can also yield the number and arrangement at the same time, and have the user do some extra unpacking, which might be a more natural approach:

import itertools

def generate_arrangements(iterable):
    iterator = iter(iterable)
    while True:
        number =  next(iterator)
        arrangement = list(list(row) for row in itertools.islice(iterator, 4))
        yield number, arrangement

(first_num, first_arr), (second_num, second_arr) = generate_arrangements(data)

As @JoranBeasley writes in the comments, this form makes it easy to use tuple unpacking in a for-loop, like this:

for num,arr in generate_arrangements(data):
    print(num) 
    print(arr)