Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Getting a Python function to cleanly return a scalar or list, depending on number of arguments

Disclaimer: I'm looking for a Python 2.6 solution, if there is one.

I'm looking for a function that returns a single value when passed a single value, or that returns a sequence when passed multiple values:

>>> a = foo(1)
2
>>> b, c = foo(2, 5)
>>> b
3
>>> c
6

To be clear, this is in an effort to make some function calls simply look nicer than:

a, = foo(1)

or

a = foo(1)[0]

Right now, the inelegant solution is something along these lines:

def foo(*args):
    results = [a + 1 for a in args]
    return results if len(results) > 1 else results[0]

Is there any syntactic sugar (or functions) that would make this feel cleaner? anything like the following?

def foo(*args):
    return *[a + 1 for a in args]
like image 636
Nevir Avatar asked Oct 27 '10 00:10

Nevir


People also ask

How do you make a function accept any number of arguments in Python?

To make a function that accepts any number of arguments, you can use the * operator and then some variable name when defining your function's arguments. This lets Python know that when that function is called with any position arguments, they should all be captured into a tuple (which that variable will point to).

How do you pass arbitrary number of arguments to a function in Python?

Well, using *args in a function is Python's way to tell that this one will: Accept an arbitrary number of arguments. Pack the received arguments in a tuple named args. Note that args is just a name and you can use anything you want instead.

Can a function have variable number of parameters in Python?

Yes. You can use *args as a non-keyword argument. You will then be able to pass any number of arguments. As you can see, Python will unpack the arguments as a single tuple with all the arguments.

How do you return multiple things from a function in Python?

In Python, you can return multiple values by simply return them separated by commas. In Python, comma-separated values are considered tuples without parentheses, except where required by syntax. For this reason, the function in the above example returns a tuple with each value as an element.


2 Answers

You can always write a decorator to elide that if statement if that is nicer to you:

import functools
def unpacked(method):
    @functools.wraps(method)
    def _decorator(*args):
        result = method(*args)
        return results if len(results) != 1 else results[0]
    return _decorator

Usage:

@unpacked
def foo(*args):
    return [arg + 1 for arg in args]
like image 107
Mike Axiak Avatar answered Sep 16 '22 13:09

Mike Axiak


You can easily write a function scalify that returns the element from the list if the list has only one element, i.e. it tries to make it a scalar (hence the name).

def scalify(l):
    return l if len(l) > 1 else l[0]

Then you can use it in your functions like so:

def foo(*args):
    return scalify([a + 1 for a in args])

This will do the trick, but I'm with those who suggest you don't do it. For one reason, it rules out iterating over the result unless you know you passed in at least two items. Also, if you have a list, you have to unpack the list when calling the function, losing its "listness," and you know you may not get a list back. These drawbacks seem to me to overshadow any benefit you may see to the technique.

like image 27
kindall Avatar answered Sep 18 '22 13:09

kindall