Specifying type for Generic Function

[hmc-cs-mdrissi]

This example comes from this comment, microsoft/pyright#2558 (comment) on ambiguity of type vars occasionally.

def to_list(value: Iterable[T] | T) -> list[T]:
    ...

def func1(x: list[int] | int):
    y = to_list(x)
    reveal_type(y)

Here when x is matched with the type of to_list there are two valid solutions. One solution is T = int, while other is T = list[int] | int. There's no rule right now on simplest type/which one is picked. I'm unsure a simplest type is even well defined/clear thing. So it'd be nice for ambiguous cases if there was a way to indicate to the generic function the intended type. Something like,

def func1(x: list[int] | int):
    y = to_list[int](x)
    reveal_type(y)

This is similar to how generic classes support __class_getitem__ as a way to indicate type specifically. One practical issue is the type annotations in a function signature are normally unused by run time. One workaround would be that functions default to not having __class_getitem__ but a decorator can be used to add one and if used type checkers would support index notation. Something like,

@generic
def to_list(value: Iterable[T] | T) -> list[T]:
    ...

def to_list2(value: Iterable[T] | T) -> list[T]:
    ... 

to_list[int] # legal at run time/type checking
to_list2[int] # Error at run time/type checking

The decorator would add a __getitem__ method to the function given that just returns the function so to_list[int] is to_list. Something like,

def generic(fn: Callable[[P], R]) -> Callable[[P], R]:
  fn.__getitem__ = lambda _: fn
  return fn

and then be used like,

def generic(fn: Callable[P, R]) -> Callable[P, R]:
    fn.__getitem__ = lambda _: fn
    return fn


@generic
def iden(x: R) -> R:
    return x


print(iden[int](1))

Testing that it fails with a type error (TypeError: 'function' object is not subscriptable), but

print(iden.__getitem__(int)(1))

works. I'm guessing __getitem__ and [] are not always same and seem to have special case for functions.

One way that currently works is to replace type ambiguous function with a generic callable object. Something like this for to_list

class ToList(Generic[T]):
  def __call__(self, x: Iterable[T] | T) -> List[T]:
    ...

ToList[int](5) # works fine

I tried making a decorator that would convert from function to a wrapper callable object but couldn't find a way to make paramspec/typevar match up to be useful. One issue is if you want to be able to index the class with the type at call sites you need generic to return Type[Wrapper] instead of Wrapper, but if you want to return Type[Wrapper] then I'm unsure how enclose the decorated function. One hacky thing that works for a single signature is,

from typing import Callable, Generic, Type, TypeVar

from typing_extensions import ParamSpec

T = TypeVar("T")


class GenericFunc(Generic[T]):
    def __init__(self, fn: Callable[[T], T]):
        self.fn = fn

    def __call__(self, x: T) -> T:
        return self.fn(x)


def generic(fn: Callable[[T], T]) -> Callable[[Type[T]], GenericFunc[T]]:
    def _generic(cls: Type[T]) -> GenericFunc[T]:
        return GenericFunc[cls](fn)

    return _generic


@generic
def iden(x: T) -> T:
    return x


reveal_type(iden(int))  # type is GenericFunc[int]

Main issue hacky solution is this works for one type variable with a specific signature. This would mean one decorator per signature type. Using paramspec doesn't seem to fit as there needs to be a way to apply just types in the signature. For a signature like,

@generic
def foo(a: int, b: bool, c: T1, d: T2) -> None:
  ...

foo[type1, type2]

would be ideal but if you try paramspec route it'd be capturing all the signature and another issue is Type[T] is a thing for typevars, but not a thing for Paramspec. For multiple type variables you'd need to adjust your generic class to have same number. Similarly unclear if pep 646 helps in this situation as I think need to know exact signature for decorator makes this approach tedious.

One way other languages (C++) handle this situation is for functions to mark their type vars. Something like,

def foo<T1, T2>(a: int, b: bool, c: T1, d: T2) -> None:
  ...


foo<T1, T2>

Adding new syntax is probably hard though so that's why I've mainly stuck to [] and not adjusting def of the function. The idea of modifying syntax with <> is partly based on these slides.

[hauntsaninja]

There is already a way to help out type checkers in ambiguous cases, use PEP 526 variable annotations for the result:

def to_list(value: Iterable[T] | T) -> list[T]:
    ...

def func1(x: list[int] | int):
    y: list[int] = to_list(x)
    reveal_type(y)

Most type checkers are capable of using this extra information during type inference, e.g. mypy calls this type context and pyright calls it bidirectional interface.

In the specific case of the to_list example, especially if you look at the proposed implementation in the linked issue, you're actually better served with overloads:

@overload
def to_list(value: Iterable[T]) -> list[T]: ...
@overload
def to_list(value: T) -> list[T]: ...

[hmc-cs-mdrissi]

Hmm, your first solution works well with adding type annotation. The second solution of overloads work with mypy, but fails with pyright and looks like a bug report to make. Specifically this code fails,

@overload
def to_list(value: Iterable[T]) -> list[T]:
    ...


@overload
def to_list(value: T) -> list[T]:
    ...


def to_list(value: Iterable[T] | T) -> list[T]: # error
    if isinstance(value, Iterable):
        return list(value)
    else:
        return [value]


def func1(x: list[int] | int):
    y: list[int] = to_list(x) # another error
    reveal_type(y)