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

Question

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]

Answer 1

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]

Answer 2

Do you mean you want a tuple with the same number of arguments? Is this not a solution?

return tuple([a + 1 for a in args])

Answer 3

def foo(*args):
    return (None, args[0]+1 if args else None, map(lambda a: a + 1, args))[len(args) if len(args) < 3 else 2]

:-) it is hell

Answer 4

This will handle 0 or more args, I think that's what you're looking for.

def foo(*args):
    return map(lambda x: x + 1, args) or [None]

edit: I revised to add a None list incase of unpacking 0 args

Answer 5

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.