Unpack a certain amount of elements in Python?

Question

Is it possible to use the splat operator (for example, a, *rest = somelist ) in such a way that it consumes a certain number of elements?

Use case: I want to break up some input that I get into a number, a list of lists, another number and another list of lists.

My input is as follows:

 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 need 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], ...] 

etc.

To do this, it would be useful to be able to consume a certain amount of elements from the iterable, I got the lines. Something like this would be ideal as an intermediate step: first_num, *[4]first_arrangement, second_num, *[4]second_arrangement = lines

What is the normal / canonical / pythonic way to solve this?

Answer 1

I think the canonical, pythonic way to do this is to put a burden on the generator that you are repeating. I would define a generator function like this:

 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 data in the list:

 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 record:

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

gives the desired result:

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

You can also specify the quantity and location at the same time, and ask the user to do additional unpacking, which can be more natural:

 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 simplifies the use of unpacking tuples in a for loop, for example:

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

Answer 2

It is not clear in what format you receive the data. Are these lines, data in a file, etc. Suppose data is a StringIO object of your input sample. You can use itertools.islice on an iterator to retrieve n elements, for example:

 >>> from itertools import islice >>> d = iter(data) >>> first_num = int(next(d)) >>> first_arrangement = [[int(x) for x in l.split()] for l in islice(d, 4)] >>> print(first_num, first_arrangement) 5 [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]