1

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?

Bruno Kim
  • 2,300
  • 4
  • 17
  • 27
  • 1
    Not exactly a duplicate, but check out discussion in https://stackoverflow.com/questions/57452652/why-does-mypy-infer-the-common-base-type-instead-of-the-union-of-all-contained-t – Samwise May 11 '22 at 16:17
  • 1
    A yucky solution would be to use `@overload` and a bunch of `one_of(arg0: T0, arg1: T1) -> Union[T0, T1]` overloads up to an arbitrary number of arguments... – Samwise May 11 '22 at 16:20

1 Answers1

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.

joel
  • 6,359
  • 2
  • 30
  • 55