Use union of arg types instead of supertype for generic variadic function

Question

I have a generic function one_of that may return any of its (variable) arguments. Its return type may be the supertype of all of its arguments; however, I'd like to annotate its usage locations with the union of their types instead.

from typing import TypeVar
import random


T = TypeVar("T")


def one_of(*args: T) -> T:
    return random.choice(args)


def int_or_str() -> int | str:
    return one_of(1, "one")

This is the error mypy provides:

vararg_type.py:11: error: Incompatible return value type (got "object", expected "Union[int, str]")
Found 1 error in 1 file (checked 1 source file)

I figured that I can cast the return value of one_of at each of its call sites, but is there another way to better use the typechecker?

Answer 1

I don't know if this is "better", but mypy only uses object because it's chosen not to default to int | str. So you could tell it to use int | str. Just hinting one is sufficient

def int_or_str() -> int | str:
    one: int | str = "one"
    return one_of(1, one)

There may be a neater way of writing that.