1

I have a few helper functions that get passed a type converter and a value. Based on checks that happen later, I decide which helper function to call.

How can I correctly annotate the types to narrow the foo variable's type below so that it can pass a mypy check?

from typing import Type, Union


def do_something(
        typ: Type[Union[float, int]],
        bar: Union[float, int]
) -> Union[float, int]:
    return bar


foo: Type[Union[float, int, str]] = float

assert foo is float or foo is int

do_something(foo, 4.4)

Bonus points if the solution can ensure that typ is a converter to the type of bar!

Martijn Pieters
  • 1,048,767
  • 296
  • 4,058
  • 3,343
Jim Hunziker
  • 14,111
  • 8
  • 58
  • 64
  • The problem isn't ``foo`` but ``typ``. The ``Type[Union[...]`` is either literally ``Union`` or undefined (since ``Union`` is not a proper class). ``typ`` should be ``Union[Type[float], Type[int], Type[str]]"``. – MisterMiyagi Mar 18 '21 at 12:41
  • That's not how the documentation shows a union of types: https://docs.python.org/3/library/typing.html?highlight=typing#typing.Type – Jim Hunziker Mar 18 '21 at 16:46

1 Answers1

3

The tool you want here is a TypeVar.

Essentially, a TypeVar let's you say "I don't know quite what type this is (although I may have some ideas), but it'll be the same one throughout its use in this function." (or in some cases throughout its use in a class)

For example, this ensures that each thing you had a Union for gets the same value on any given call to the function.

from typing import Type, TypeVar

# Define a type variable
# and list the things it is allowed to represent
NumberType = TypeVar("NumberType", int, float) 

def do_something(
        typ: Type[NumberType],
        bar: NumberType
) -> NumberType:
    return bar

This can be legally called with do_something(float, 2.5) in which case it will return a float, or it can be called with do_something(int, 2) in which case it will return an int. That is, it makes sure that all the things match.

Because you called it a type converter, I suspect that you may not actually want all the types to match. If you need to constrain on more than one Type Variable, you can do so with something more like

from typing import Callable, TypeVar

# Define a type variable
# and list the things it is allowed to represent
NumberTypeIn = TypeVar("NumberTypeIn", int, float)
NumberTypeOut = TypeVar("NumberTypeOut", int, float) 

def do_something(
        converter: Callable[[NumberTypeIn], NumberTypeOut],
        bar: NumberTypeIn
) -> NumberTypeOut:
    return type_(bar)

As to the original question of narrowing unions of Type[]s, as you've noticed is doesn't work. Instead you can use issubclass, as in

assert not issubclass(foo, str)

or

assert issubclass(foo, int) or issubclass(foo, float)
Josiah
  • 1,072
  • 6
  • 15
  • 1
    Also, this is not what you're asking, but it may be relevant to mention because it'll cause confusion otherwise. As far as MyPy is concerned, an `int` is a `float` because it duck types it. (https://mypy.readthedocs.io/en/stable/duck_type_compatibility.html) If you want to play around with this, it may be easier to use less confusing types like str and dict. – Josiah Mar 18 '21 at 23:37
  • `foo` is no longer in your example. Will this work if `foo` is a variable that's typed to a union of three things while `do_something` is typed to two things like in my post? – Jim Hunziker Mar 19 '21 at 12:45
  • Also, I called it a converter because the constructor(?) of `float` and `int` converts strings to those types, but really it will only be `float`, `int`, or `str`. I don't know if those have a supertype in Python. – Jim Hunziker Mar 19 '21 at 12:49
  • Ah. Sorry, I missed that bit. I think you can use `issubclass` in your assert. E.g. `assert not issubclass(foo, str)` or `assert issubclass(foo, int) or issubclass(foo, float)` narrow things correctly. – Josiah Mar 19 '21 at 14:47
  • 1
    Excellent! Can you put that in your answer so I can mark it as accepted? – Jim Hunziker Mar 19 '21 at 20:34