How to annotate this function that wraps non-tuples/lists in a tuple and returns tuples & lists unchanged?

[jwodder]

Consider the following function:

def ensure_tuple_or_list(obj):  
    if isinstance(obj, (list, tuple)):
        return obj
    return (obj,)

How exactly do I annotate this better than the information-lossy (obj: Any) -> list | tuple?

My first attempt:

from typing import TypeVar, overload

T = TypeVar("T")

@overload
def ensure_tuple_or_list(obj: list[T]) -> list[T]:
    ...

@overload
def ensure_tuple_or_list(obj: tuple) -> tuple:
    ...

@overload
def ensure_tuple_or_list(obj: T) -> tuple[T]:
    ...

def ensure_tuple_or_list(obj: list[T] | tuple | T) -> list[T] | tuple:
    if isinstance(obj, (list, tuple)):
        return obj
    return (obj,)

Mypy 1.1.1 with default settings rejects this with:

ensure-tuple-or-list01.py:6: error: Overloaded function signatures 1 and 3 overlap with incompatible return types  [misc]
ensure-tuple-or-list01.py:10: error: Overloaded function signatures 2 and 3 overlap with incompatible return types  [misc]
Found 2 errors in 1 file (checked 1 source file)

---

For my second attempt, I tried having the function always return a tuple:

from typing import TypeVar, overload

T = TypeVar("T")

@overload
def ensure_tuple_or_list(obj: list[T]) -> tuple[T, ...]:
    ...

@overload
def ensure_tuple_or_list(obj: tuple) -> tuple:
    ...

@overload
def ensure_tuple_or_list(obj: T) -> tuple[T]:
    ...

def ensure_tuple_or_list(obj: list[T] | tuple | T) -> tuple:
    if isinstance(obj, (list, tuple)):
        return tuple(obj)
    return (obj,)

Mypy rejects this with:

ensure-tuple-or-list02.py:6: error: Overloaded function signatures 1 and 3 overlap with incompatible return types  [misc]
ensure-tuple-or-list02.py:10: error: Overloaded function signatures 2 and 3 overlap with incompatible return types  [misc]
ensure-tuple-or-list02.py:17: error: Overloaded function implementation does not accept all possible arguments of signature 1  [misc]
Found 3 errors in 1 file (checked 1 source file)

---

For my third attempt, I switched to using functools.singledispatch:

from functools import singledispatch
from typing import TypeVar

T = TypeVar("T")

@singledispatch
def ensure_tuple_or_list(obj: T) -> tuple[T]:
    return (obj,)

@ensure_tuple_or_list.register
def _(obj: list[T]) -> list[T]:
    return obj

@ensure_tuple_or_list.register
def _(obj: tuple) -> tuple:
    return obj

Mypy rejects this with:

ensure-tuple-or-list03.py:10: error: Argument 1 to "register" of "_SingleDispatchCallable" has incompatible type "Callable[[List[T]], List[T]]"; expected "Callable[..., Tuple[T]]"  [arg-type]
ensure-tuple-or-list03.py:11: error: Dispatch type "List[T]" must be subtype of fallback function first argument "T"  [misc]
ensure-tuple-or-list03.py:14: error: Argument 1 to "register" of "_SingleDispatchCallable" has incompatible type "Callable[[Tuple[Any, ...]], Tuple[Any, ...]]"; expected "Callable[..., Tuple[T]]"  [arg-type]
ensure-tuple-or-list03.py:15: error: Dispatch type "Tuple[Any, ...]" must be subtype of fallback function first argument "T"  [misc]
Found 4 errors in 1 file (checked 1 source file)

Variations on this last one (which I will omit for brevity) similarly proved fruitless.

---

Help?

[gandhis1]

This is kind of similar to the example here, although without respect to overloads: https://github.com/microsoft/pyright/blob/main/docs/mypy-comparison.md#constraint-solver-ambiguous-solution-scoring

Curious, what do you get from Pyright?

[jwodder]

I do not use pyright, and I have no intention of using pyright. Please only give answers applicable to mypy.

[Akuli]

Your first attempt "works", if I change tuple to tuple[T, ...] in a few places and I add a few #type: ignores.

from typing import TypeVar, overload, Any, Union

T = TypeVar("T")

@overload
def ensure_tuple_or_list(obj: list[T]) -> list[T]:  # type: ignore
    ...

@overload
def ensure_tuple_or_list(obj: tuple[T, ...]) -> tuple[T, ...]:  # type: ignore
    ...

@overload
def ensure_tuple_or_list(obj: T) -> tuple[T]:
    ...

def ensure_tuple_or_list(obj: list[T] | tuple[T, ...] | T) -> list[T] | tuple[T, ...]:
    if isinstance(obj, (list, tuple)):
        return obj  # type: ignore
    return (obj,)


reveal_type(ensure_tuple_or_list(1))  # Tuple[int]
reveal_type(ensure_tuple_or_list((1,)))  # Tuple[int]
reveal_type(ensure_tuple_or_list([1]))  # list[int]

You need to use # type: ignore (or some other way of hiding errors) to make mypy happy with this. To some extent, it is by design. The type of ensure_tuple_or_list(foo) depends on whether foo is known to be a list or tuple. If it is of an unknown type, then mypy has no way to know whether your function will wrap it or not, and you will get some unexpected results.

Consider this, for example:

x: int | list[int]
reveal_type(ensure_tuple_or_list(x))  # Tuple[Union[builtins.int, builtins.list[builtins.int]]]

Because the type int | list[int] is not known to be a list or a tuple, mypy decides to wrap it in a tuple, and you get back an annoying tuple[int | list[int]]. There doesn't seem to be a good workaround that would give you the expected tuple[int] | list[int] (I tried a couple things). So even though it's hard to type this function in a way that works in all cases, your first idea kind of works with small changes.