Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python type hints for type promotion

Consider a function that performs type promotion, e.g. a simple multiplication of two numbers that can both be either int or float:

def mul(a: int | float, b: int | float):  # return type?
    return a * b

This function returns float, except in the case where both a and b are int.

How can I properly and concisely annotate the return type? I know I can do this with @overload:

from typing import overload

@overload
def mul(a: int, b: int) -> int: ...

@overload
def mul(a: float, b: int | float) -> float: ...

@overload
def mul(a: int | float, b: float) -> float: ...

def mul(a, b):
    return a * b

but this is very verbose and requires many overloads for something I would imagine some "type function" should handle. In C++ this could be done e.g. with SFINAE. Is there something similar I can do in Python in terms of a generic function along the lines of

def mul(a: T1, b: T2) -> promote_types(T1, T2):
    return a * b

that also works with TypeVars? I don't expect anything built in that already works for int and float, but some technique perhaps?

Notes:

  • I know about the recommendation to just annotate everything taking an int with float, but my setting has more complicated TypeVars, the choice of int and float here is just a simple example.

  • I know I can just do Union[int, float], but I need it to be specific. Depending on the exact types the function is called with, the return type must be exact too, not a union.

like image 233
Darkdragon84 Avatar asked Nov 24 '25 17:11

Darkdragon84


1 Answers

In my understanding of the docs about Generics, the following should work:

from typing import TypeVar, Generic

T1 = TypeVar('T1', str, float)
T2 = TypeVar('T2', str, float)

class Promoted(Generic[T1, T2]):
    def __class_getitem__(cls, types):
        t1, t2 = types
        if t1 == str or t2 == str:
            return str
        return float

def method(a: T1, b: T2) -> Promoted[T1, T2]:
    if isinstance(a, str) or isinstance(b, str):
        return f"{a} {b}"
    else:
        return a*b

result1 = method("Hello", "World")  # should be inferred as str
result2 = method(3.0, 4.0)  # should be inferred as float
result3 = method("Price", 5.0)  # should be inferred as str

but it does not (at least in my IDE, IntelliJ). In particular, all results are type-hinted as "Promoted", instead of str or float. Maybe you get some more interesting behavior from your inspector?

like image 155
Gabriele Castellano Avatar answered Nov 26 '25 06:11

Gabriele Castellano