Are type vars and type unions incompatible in python?

Question

Consider the following snippet:

from typing import TypeVar

import numpy as np

T = TypeVar("T", float, np.ndarray)


def f(x: T) -> T:
    """
    expects a float or an array and returns an output of the same type
    """
    return x * 2


f(1)  # ok

f(np.array([1, 2, 3]))  # ok


def g(x: float | np.ndarray) -> float | np.ndarray:
    """
    expects either a float or an array
    """
    return f(x) / 2  # should be fine, but pyright complains about type

I have created a TypeVar to hint that f expects as input a float or an array and will return an output of the same type.

The type hint in g is more loose. It expects either a float or an array and will return a float or an array, without constraining the type of the output to the type of the input.

Intuitively, the setup makes sense. Inside the definition of the g function we know that we expect x to be either a float or an array, i.e. what f expects as input. However when I pass x to f at the last line, Pyright complains:

Argument of type "float | ndarray[Unknown, Unknown]" cannot be assigned to parameter "x" of type "T@f" in function "f"

Type "float | ndarray[Unknown, Unknown]" is incompatible with constrained type variable "T"

This is surprising and frustrating, because it means that one cannot use my function f without being very cautious about the way they write their type hints.

Any thoughts on how to solve this?

Edit: After the comment of Brian61354270, I have recreated essentially the same example, only with no dependence of numpy. Here instead of numpy array we use Fraction:

from fractions import Fraction
from typing import TypeVar

T = TypeVar("T", float, Fraction)


def f(x: T) -> T:
    """
    expects a float or a Fraction and returns an output of the same type
    """
    return x * 2


f(1.0)  # ok

f(Fraction(1, 2))  # ok


def g(x: float | Fraction) -> float | Fraction:
    """
    expects either a float or a Fraction
    """
    return f(x) / 2  # should be fine, but pyright complains about type

Again, Pyright reports essentially the same issue:

Argument of type "float | Fraction" cannot be assigned to parameter "x" of type "T@f" in function "f"

Type "float | Fraction" is incompatible with constrained type variable "T"

Interestingly, if instead of Fraction we use int, the type check passes:

from typing import TypeVar

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


def f(x: T) -> T:
    """
    expects a float or an integer and returns an output of the same type
    """
    return x * 2


f(1.0)  # ok

f(1)  # ok


def g(x: float | int) -> float | int:
    """
    expects either a float or an integer
    """
    return f(x) / 2  # now its ok

Answer 1

Are type vars and type unions incompatible in python?

This isn't the issue. The type variable defined, T = TypeVar("T", float, Fraction), is constrained to float or Fraction. f(x: T) can only accept an instance of float or Fraction, not the union float | Fraction. Your example can be reduced to the following (Pyright Playground, mypy Playground)

from typing import TypeVar
from fractions import Fraction

T = TypeVar("T", float, Fraction)

def f(x: T) -> T: ...
def getFloatOrFraction() -> float | Fraction: ...

>>> f(1.0)  # ok
>>> f(Fraction(1, 2))  # ok
>>> num: float | Fraction = getFloatOrFraction()
>>> f(num)  # Error

I don't exactly know what your f and g needs, but

• If you still require type constraints (e.g. if all the inferred return type of f should be upcasted to exactly one of float, Fraction, or float | Fraction), you'll need to add the union as an additional constraint:

  T = TypeVar("T", float, Fraction, float | Fraction)
  
  class MyFloat(float): ...
  
  def getMyFloatOrFraction() -> MyFloat | Fraction: ...
  

  >>> reveal_type(f(getMyFloatOrFraction()))  # float | Fraction
  

• If you don't require type constraints, then setting an upper bound is easier:

  T = TypeVar("T", bound=float | Fraction)
  
  class MyFloat(float): ...
  
  def getMyFloatOrFraction() -> MyFloat | Fraction: ...
  

  >>> reveal_type(f(getMyFloatOrFraction()))  # MyFloat | Fraction
  

Answer 2

While not a desirable solution, you can resolve this error by repeating the call to f in g under explicit type guards:

def g(x: float | np.ndarray) -> float | np.ndarray:
    """
    expects either a float or an array
    """
    if isinstance(x, (float, int)):
        return f(x) / 2
    return f(x) / 2

Output using pyright v1.1.329:

0 errors, 0 warnings, 0 informations