Union math is too greedy

[ilevkivskyi]

Currently union math can produce excessively wide types for overload calls. Example:

from typing import Union, overload

@overload
def f(x: Union[int, str]) -> int: ...
@overload
def f(x: object) -> object: ...
def f(x):
    pass

x: Union[int, str]
reveal_type(f(x))  # object???

A possible solution would be to check if the "normal way" generates a better result ad return it instead of returning soon. But maybe we should actually re-think the union math algorithm. I will spend some time thinking about it, and will make a PR if I will find an easy solution.

[Michael0x2a]

Another possible solution would be to take the call argument types, and produce multiple copies of where we vary any arguments that have a union type. We then take each copy and try finding the first match then union those return types together.

So here, we'd try to match f(int) and f(str) then union the return types together to get Union[int, int], which simplifies to the expected result.

This algorithm would also work with the example that convinced me to do union math first, rather then second.

The only downside is that this would end up being pretty expensive in the case where the call contains many args with unions, or very large unions. I guess it probably wouldn't be too hard to add in some checks to bail out if union math looks expensive though.

[ilevkivskyi]

The only downside is that this would end up being pretty expensive in the case where the call contains many args with unions, or very large unions.

Yes. Choosing best from unioned and normal is also a bit slower than currently. What you propose is exactly what I was thinking about, but it may be actually quadratic in worst case and will only buy us more precise types in some corner cases. I therefore think we should go with the simpler solution for now, if someone will ask we could invent something more sophisticated.

I already have branch ready, will make a PR in few minutes.

[Michael0x2a]

but it may be actually quadratic in worst case and will only buy us more precise types in some corner cases.

Yeah -- I was thinking we could probably steal some of the logic from the current implementation which effectively prevents union math from happening if two or more arguments end up potentially varying. (E.g. we currently mandate the unioned result must have only a single union argument).

We could borrow that same idea to make sure we only destructure a single union argument to avoid entering into the quadratic case (even if the call contains multiple unions) -- the cost would be that we'd have to do an extra linear scan or two.

But in any case, I'm also ok with just going with the simpler solution for now. Union math is supposed to be a convenience, anyways -- if the user is doing something more complicated, the workarounds are easy enough.