Is it good code style if a Python function returns different types depending on the arguments

Question

Is it good code style if a Python function returns different types depending on the arguments provided?

def foo(bar):
  if bar is None:
    return None
  elif bar == 1:
    return 1*1
  else:
    return [b*b for b in bar]

foo returns None if bar is None

foo return 1 if bar == 1

foo returns a List of int if bar is a Tuple / List of integers

Examples:

>> foo(None)
None
>> foo(1)
1
>> foo(1, 2, 3, 4)
[1, 4, 9, 16]

Returning None or an int should be OK, but is it OK to return an int or a list of ints depending on the function arguments? You could argue that it would be OK because the user knows which return types to expect and doesn't need type checking (in which case I would say it's not OK), but one could argue that it would be better to split the function into two functions, one expecting a int and return an int and one expecting a list of int and returning a list of ints.

Answer 1

It depends entirely on your use case. Here's an example in the standard library where the type of the result depends on the type of the input:

>>> import operator
>>> operator.add(1, 2)
3
>>> operator.add(1.0, 2.0)
3.0

This type of behavior is usually ok, and can be documented with @typing.overload.

---

Here's an example where the type of the result depends on the value of the input:

>>> import json
>>> json.loads('1')
1
>>> json.loads('[1]')
[1]

This type of behavior is usually reserved for serialization and deserialization, or APIs which blur the type / value boundary such as astype in np.int_(3).astype(bool).

---

On the other hand, here's an example of a function which is obviously poorly designed:

from typing import Union  # make sure to document the mixed return type

def is_even(x: int) -> Union[bool, str]:
    if x % 2 == 0:
        return True
    else:
        return "no"

Without knowing your specific use case, it's hard to give advice here.