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

Question

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.

Answer 1

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.