Search code examples
pythongenericspython-typingmypy

How can I have an optional TypeVar in a Generic class in Python?


I'm trying to write a simple type wrapper to represent the interface of decorator functions:

from typing import Protocol, TypeVar, Generic

TIn = TypeVar('TIn', contravariant=True)
TOut = TypeVar('TOut', covariant=True)

class Decorator(Protocol, Generic[TIn, TOut]):
    """
    Represents a decorated value, used to simplify type definitions
    """
    def __call__(self, value: TIn) -> TOut:
        ...

This would be used to type a decorator function as follows:

IntFunction = Callable[[int, int], int]


def register_operator(op: str) -> Decorator[IntFunction, IntFunction]:
    def inner(value: IntFunction) -> IntFunction:
        # register the function or whatever
        return value
    return inner


@register_operator("+")
def add(a: int, b: int) -> int:
    return a + b

In the above example, Mypy is able to validate the type signature of add to ensure it matches the specification of register_operator.

This is useful for decorators that transform the type (eg converting it from an IntFunction to a StrFunction), but in almost all cases, TIn is identical to TOut, and so I want to simplify the usage of my definition.

Essentially, I want to make it so that if TOut isn't given, it will be assumed to be the same as TIn, which would allow the above decorator function to be simplified to

def register_operator(op: str) -> Decorator[IntFunction]:
    #                              Simplification here ^

    def inner(value: IntFunction) -> IntFunction:
        # register the function or whatever
        return value
    return inner

The ideal syntax I would use in my protocol definition would be something like this:

class Decorator(Protocol, Generic[TIn, TOut = TIn]):
    """
    Represents a decorated value, used to simplify type definitions
    """
    def __call__(self, value: TIn) -> TOut:
        ...

Note that this does not work.

How can I achieve this functionality, whilst continuing to have the assurance that Mypy provides? I am happy to make the definition of Decorator as complex as needed, but I want to keep its simple usage.


Solution

  • So as far I know python doesn't support default values in generics, a generic that isn't fully defined becomes unknown. Perhaps it should support it, there's a draft pep for this feature (see update below).

    For now, I would create two classes that expose different interfaces. This is almost as good, it just means you need a second class:

    from typing import Protocol, TypeVar, Generic, overload
    
    TIn = TypeVar('TIn', contravariant=True)
    TOut = TypeVar('TOut', covariant=True)
    TSym = TypeVar('TSym')
    
    class Decorator(Protocol, Generic[TIn, TOut]):
        """
        Represents a decorated value, used to simplify type definitions
        """
        def __call__(self, value: TIn) -> TOut:
            ...
    
    class SymmetricDecorator(Decorator[TSym,TSym], Generic[TSym], Protocol):
        pass
    

    By "Symmetric" I just mean the callable receives and returns the same type - there might be a better name. Hope this is useful!

    UPDATE: I've since found there exists an active pep for default values for TypeVar: https://peps.python.org/pep-0696/