Controlling the number of returned variables in a function

Question

Is there a way to define a function which knows how many variables to return based on how many outputs the user expects?

Let me illustrate the idea. Assume the following function:

def function():
    a = 1
    b = 2
    c = 3
    return a, b, c

Then, I would like my function to behave like:

>>> x = function()
>>> x
1

>>> x, y = function()
>>> x
1
>>> y
2

>>> x, y, z = function()
>>> x
1
>>> y
2
>>> z
3

Is there a function or concept in python that can help me to achieve that? Maybe decorators?

Any idea is welcome!

PD: My level of Python is still very basic.

---

EDIT:

I am currently moving from IDL to Python. So I am missing the nice feature in IDL that you can actually choose to return as many variables as desired by doing like:

FUNCTION function, a=a, b=b, c=c

 a=1
 b=2
 c=3

RETURN, a

And then you can simply ask what you want to get back

IDL> x=function()
IDL> print, x
1
IDL> x=function(y=b)
IDL> print, x
1
IDL> print, y
2
IDL> x=function(y=b, z=c)
IDL> print, x
1
IDL> print, y
2
IDL> print, c
3

Answer 1

Instead of returning different values, you could call it differently:

x, *_ = function()
x, y, *_ = function()
x, y, z, *_ = function()   # *_ is optional here if it only returns 3 things

Doing this assigns all unused returned values to the _ variable, so if you want them to get gc'd early, you have to del _.

Answer 2

You can only ever return a single object. Note, your function is returning a tuple. However, Python supports syntax that lets you unpack variables flexibly

x,_,_ = function()

Or

x,y,_ = function()

Or even with extended unpacking:

x, *_ = function()

Note, using _ as a throwaway variable is merely a convention.

Answer 3

Short of insane tricks like examining your caller’s bytecode to see how many values it’s expecting, this is impossible: the interpreter will not accept any mismatched number of values (it checks for and traps extras).

Answer 4

While not as Pythonic as unpacking, you can create a decorator and specify the number of returned values that are allowed:

def filter_results(val):
  def outer(f):
    def wrapper(*args, **kwargs):
       v = f(*args, **kwargs)[:val]
       return v[0] if len(v) == 1 else v
    return wrapper
  return outer

@filter_results(1)
def function():
  a = 1
  b = 2
  c = 3
  return a, b, c

r = function()

Output:

1

Full results:

@filter_results(2)
def function():
   ...

x, y = function()

Lastly:

@filter_results(3)
def function():
   ...

x, y, z = function()