python typing Dict and overloading: Overloaded function implementation does not accept all possible arguments

Question

@overload
def get_random(
        d:Dict[int,int]
)->int:
    ...

@overload
def get_random(
        d:Dict[int,float]
)->float:
    ...

def get_random(
        d: Dict[int,Union[int,float]]
)->Union[int,float]:
    key = random.choice(list(d.keys()))
    return d[key]


a: int = get_random({1:1})
b: float = get_random({1:1.})

mypy provides the following message:

error: Overloaded function implementation does not accept all possible arguments of signature 1

I am afraid I do not understand what the issue is.

(note: I am aware that an alternative working solution is to use TypeVar)

Answer 1

Note that Dict[int, Union[int, float]] is not the same as Union[Dict[int, int], Dict[int, float]] which is what you meant.

import random
from typing import Dict, Union, overload


@overload
def get_random(d: Dict[int, int]) -> int:
    ...


@overload
def get_random(d: Dict[int, float]) -> float:
    ...


def get_random(d: Union[Dict[int, int], Dict[int, float]]) -> Union[int, float]:
    key = random.choice(list(d.keys()))
    return d[key]


reveal_type(get_random({1: 1.0}))  # Revealed type is 'builtins.float'
reveal_type(get_random({1: 1}))  # Revealed type is 'builtins.int'

mypy still gives an error here due to its strict policy about overlapping overloaded variants. Note that int is a subtype of float, therefore, for example, {1: 2} will match both variants. Read more here.

If the order of variants is swapped, the discrepancy is clear, as the revealed type is float in both cases.

import random
from typing import Dict, Union, overload


@overload
def get_random(d: Dict[int, float]) -> float:
    ...


@overload
def get_random(d: Dict[int, int]) -> int:
    ...


def get_random(d: Union[Dict[int, int], Dict[int, float]]) -> Union[int, float]:
    key = random.choice(list(d.keys()))
    return d[key]


reveal_type(get_random({1: 1.0}))  # Revealed type is 'builtins.float'
reveal_type(get_random({1: 1}))  # Revealed type is 'builtins.float'

Consider using a TypeVar instead.

import random
from typing import Dict, TypeVar

Number = TypeVar("Number", int, float)


def get_random(d: Dict[int, Number]) -> Number:
    key = random.choice(list(d.keys()))
    return d[key]


reveal_type(get_random({1: 1.0}))  # Revealed type is 'builtins.float'
reveal_type(get_random({1: 1}))  # Revealed type is 'builtins.int'