Returning unpacked list in python

Viewed 1167

I'm trying to create a function in python that takes any number of arguments, runs some operations on the arguments and lastly returns the new values. Eg:

def foo(*args):
    return *[i ** 2 for i in args] # error: cant use starred expression here

a, b, c, d = foo(2, 15, 17, 21)

instead of "unpacking" the list it just spits out a syntax error. I could technically leave out the "unpacking" all together and just have the user write

def foo(*args):
    return [i ** 2 for i in args]

a, b, c, d = foo(2, 15, 17, 21) # works

e = foo(82) # e is now a list and not just the value

this works when having multiple arguments but not when I only give one. Is there any way to do this easily without the need of using an if-statement to check the number of inputs given?

2 Answers

You can unpack a one value list, you just need a comma after the one name you're assigning to:

e, = foo(82)

This is the same kind of syntax you need to use when creating a 1-tuple. You need to use (x,), because (x) is identical to x (the parentheses don't create the tuple, the comma does).

def foo(*args):
    return [i ** 2 for i in args]

a, b, c, d = foo(2, 15, 17, 21)

works as intended. Of course, you must ensure that the returned list has exactly as many items as you have on the left side of the assignment.

Related