23

What is the best way in Python to automatically extend a list to N elements, if the list has fewer than N elements?

That is, let's have I have this string: s = "hello there". If I do this:

x, y, z = s.split()

I will get an error, because s.split() returns a list of two elements, but I'm assigning it to 3 variables. What I want is for z to be assigned None.

I know I can do this the hard way:

l = s.split()
while len(l) < 3:
    l.append(None)
x, y, z = l

But there has to be something more elegant than this.

charlesreid1
  • 4,360
  • 4
  • 30
  • 52
Richard Fallon
  • 537
  • 4
  • 8

7 Answers7

26

If you want a one-liner, do:

s = "hello there"
x, y, z, *_ = s.split() + [None, None, None]

print(x, y, z)

Output

hello there None

Note that a one-liner is not necessarily more readable nor more elegant. A variation, thanks to @Grismar is:

x, y, z, *_ = s.split() + [None] * 3
Dani Mesejo
  • 61,499
  • 6
  • 49
  • 76
  • 7
    Depending on the number of `None`s, I'd prefer `x, y, z, *_ = s.split() + [None] * 3` - it has the number right in there, which would be more readable for a larger number, makes it easier to parameterise and is shorter to boot. – Grismar Dec 02 '19 at 00:20
  • 4
    If you want to get rid of `*_`, you can write it `x, y, z = (s.split() + [None]*3)[:3]` – Andrej Kesely Dec 02 '19 at 01:06
  • 1
    itertools.chain(s.split(), itertools.repeat(None)) – RiaD Dec 02 '19 at 12:39
  • 1
    @RiaD: Good idea, but with `x, y, z = chain(s.split(), repeat(None)` you'll get a `ValueError: too many values to unpack`, and with `x, y, z, *_ = chain(s.split(), repeat(None)` you'll get an attempt to build an infinite list in `_` which will block until it has eaten all available memory. – Seb Dec 02 '19 at 15:41
  • ah, you are right, so have to add slice and it's becoming too long – RiaD Dec 02 '19 at 16:31
16

extend adds an iterable to the end of a list, so you can do:

l.extend(None for _ in range(3 - len(l))

or 

l.extend([None]*(3-len(l)))

which is a bit more elegant but slightly slower because it needs to construct the list of Nones first.

geisterfurz007
  • 5,292
  • 5
  • 33
  • 54
Steven Fontanella
  • 764
  • 1
  • 4
  • 16
7

This sort of application is exactly what the padnone recipe in itertools is for (https://docs.python.org/3/library/itertools.html#itertools-recipes):

x, y, z = itertools.islice(itertools.chain(s.split(), itertools.repeat(None)), 3)

The above also incorporates take to get just 3 elements.

Iguanodon
  • 71
  • 1
  • 3
  • Be careful not to run the variant to this without islice: `x, y, z, *_ = itertools.chain(s.split(), itertools.repeat(None))` as it will happily try and create a infinite iterable in `_`. Providing a value for amount of repetitions works though `x, y, z, *_ = itertools.chain(s.split(), itertools.repeat(None, 10))` – user1556435 Jul 29 '21 at 13:05
6

Here's a mildly amusing solution using itertools.zip_longest: the behaviour is to use None to fill in the missing elements in the shorter sequence.

from itertools import zip_longest

(x, _), (y, _), (z, _) = zip_longest(s.split(), range(3))

Like your original code, this will throw an error if s.split() has more than three parts. If you prefer to silently discard the extra parts, destructure to (x, _), (y, _), (z, _), *_ instead.

kaya3
  • 47,440
  • 4
  • 68
  • 97
4

An alternative approach would be to assign z to a blank string using partition. 

s = "hello there" 
x, z, y = s.partition(' ')

This approach cuts to the chase. In doing so, it leaves z assigned to a blank instead of None.

gregory
  • 10,969
  • 2
  • 30
  • 42
3

A flexible decorator-based solution inspired by this and this answer:

from itertools import islice, chain, repeat

def variable_return(max_values, default=None):
    def decorator(f):
        def wrapper(*args, **kwargs):
            return islice(chain(f(*args, **kwargs), repeat(default)), max_values)
        return wrapper
    return decorator

Or a pointlessly condensed version:

variable_return = lambda m, d=None: lambda f: lambda *a, **k: islice(chain(f(*a, **k), repeat(d)), m)

It can be used like this:

@variable_return(3)
def split(s):
    return s.split()

x, y, z = split(s) # ('hello', 'there', None)

or like this:

x, y, z = variable_return(3)(s.split)() # ('hello', 'there', None)
Seb
  • 4,422
  • 14
  • 23
3

You could add the single-item list [None] the required number of times (edit: similar to this but without using extend):

l += [None]*(3-len(l))
Charley
  • 35
  • 4