CODEWITHSUNDEEP
pythonpython-3.x
Puff
Puff

Reputation: 515

Python: is it possible to return a variable amount of variables as requested?

I'd like to return either one or two variables for a function in python(3.x). Ideally, that would depend on amount of returned variables requested by user on function call. For example, the max() function returns by default the max value and can return the argmax. In python, that would look something like:

maximum = max(array)
maximum, index = max(array)

Im currently solving this with an extra argument return_arg:

import numpy as np

def my_max(array, return_arg=False):
   maximum = np.max(array)

   if return_arg:
      index = np.argmax(array)
      return maximum, index
   else:
      return maximum

This way, the first block of code would look like this:

maximum = my_max(array)
maximum, index = my_max(array, return_arg=True)

Is there a way to avoid the extra argument? Maybe testing for how many vaules the user is expecting? I know you can return a tuple and unpack it when calling it (that's what I'm doing).

Asume the actual function I'm doing this in is one where this behaviour makes sense.

Upvotes: 5

Views: 5062

Answers (4)

blhsing
blhsing

Reputation: 107095

If you need a solution that works for any data type such as np.ndarray, you can use a decorator that uses ast.NodeTransformer to modify any assignment statement that assigns a call to a given target function name (e.g. my_max) to a single variable name, to the same statement but assigns to a tuple of the same variable name plus a _ variable (which by convention stores a discarded value), so that a statement such as maximum = my_max(array) is automatically transformed into maximum, _ = my_max(array):

import ast
import inspect
from textwrap import dedent

class ForceUnpack(ast.NodeTransformer):
    def __init__(self, target_name):
        self.target = ast.dump(ast.parse(target_name, mode='eval').body)
    def visit_Assign(self, node):
        if isinstance(node.value, ast.Call) and ast.dump(node.value.func) == self.target and isinstance(node.targets[0], ast.Name):
            node.targets[0] = ast.Tuple(elts=[node.targets[0], ast.Name(id='_', ctx=ast.Store())], ctx=ast.Store())
            return node
    # remove force_unpack from the decorator list to avoid re-decorating during exec
    def visit_FunctionDef(self, node):
        node.decorator_list = [
            decorator for decorator in node.decorator_list
            if not isinstance(decorator, ast.Call) or decorator.func.id != "force_unpack"
        ]
        self.generic_visit(node)
        return node

def force_unpack(target_name):
    def decorator(func):
        tree = ForceUnpack(target_name).visit(ast.parse(dedent(inspect.getsource(func))))
        ast.fix_missing_locations(tree)
        scope = {}
        exec(compile(tree, inspect.getfile(func), "exec"), func.__globals__, scope)
        return scope[func.__name__]
    return decorator

so that you can define your my_max function to always return a tuple:

def my_max(array, return_arg=False):
   maximum = np.max(array)
   index = np.argmax(array)
   return maximum, index

while applying the force_unpack decorator to any function that calls my_max so that the assignment statements within can unpack the returning values of my_max even if they're assigned to a single variable:

@force_unpack('my_max')
def foo():
    maximum = my_max(array)
    maximum, index = my_max(array)

Upvotes: 1

Kurtis Streutker
Kurtis Streutker

Reputation: 1316

No, there is no way of doing this. my_max(array) will be called and return before assigning a value to maximum. If more than one value is returned by the function then it will try unpacking the values and assigning them accordingly.
Most people tackle this problem by doing this:

maximum, _ = my_max(array)
maximum, index = my_max(array)

or

maximum = my_max(array)[0]
maximum, index = my_max(array)

Upvotes: 2

blhsing
blhsing

Reputation: 107095

You can instead return an instance of a subclass of int (or float, depending on the data type you want to process) that has an additional index attribute and would return an iterator of two items when used in a sequence context:

class int_with_index(int):
    def __new__(cls, value, index):
        return super(int_with_index, cls).__new__(cls, value)
    def __init__(self, value, index):
        super().__init__()
        self.index = index
    def __iter__(self):
        return iter((self, self.index))

def my_max(array, return_arg=False):
   maximum = np.max(array)
   index = np.argmax(array)
   return int_with_index(maximum, index)

so that:

maximum = my_max(array)
maximum, index = my_max(array)

would both work as intended.

Upvotes: 2

Hugo G
Hugo G

Reputation: 16536

The answer is no, in Python a function has no context of the caller and can't know how many values the caller expects in return.

Instead in Python you would rather have different functions, a flag in the function signature (like you did) or you would return an object with multiple fields of which the consumer can take whatever it needs.

Upvotes: 2

Related Questions