Search code examples
python-3.xtypesmypy

Is it possible in python to specify that a generic type has some attributes and methods


In python I would like to create a function that would accept a generic parameter and that would return a value of the same type and I would like to somehow specify that the generic has some attributes and methods. I thought that binding a TypeVar to a Protocol would do the job, but I was terribly wrong

Just to give an example of such a function, let's say the function sorts the input array of some objects and returns sorted array of those objects. The only restriction on the array's elements is that they have to implement __lt__

from __future__ import annotations
from typing import List, Protocol, TypeVar

class ComparableProto(Protocol):
    def __lt__(self, other: ComparableProto) -> bool: ...

Comparable = TypeVar("Comparable", bound=ComparableProto)

def sort(arr: List[Comparable]) -> List[Comparable]: ...

if __name__ == "__main__":
    sort([1, 2, 3])  # mypy will say: Value of type variable "Comparable" of "sortarr" cannot be "int"  [type-var]

I understand that the above example can't work, because Comparable type must be subtype of ComparableProto and just because some type implements __lt__ doesn't make it subtype of ComparableProto.

I noticed that typing module implements several generic protocol types named Supports*, but they differ from what I want

Is there some way to achieve what I want?


Solution

  • Your code is almost fine. There's nothing wrong with TypeVar bound to Protocol, you use it exactly as expected. The problem is in protocol definition: int is not compatible with ComparableProto, because it has def __lt__(self, __other: int) -> bool. Notice that int is comparable only with another int, while your protocol expects wider type, ComparableProto. You can verify that with a simpler program:

    from __future__ import annotations
    from typing import Protocol
    
    class ComparableProto(Protocol):
        def __lt__(self, other: ComparableProto) -> bool: ...
    
    x: ComparableProto = 1
    

    mypy will explain, what's happening, with notes:

    main.py:7: error: Incompatible types in assignment (expression has type "int", variable has type "ComparableProto")
    main.py:7: note: Following member(s) of "int" have conflicts:
    main.py:7: note:     Expected:
    main.py:7: note:         def __lt__(self, ComparableProto) -> bool
    main.py:7: note:     Got:
    main.py:7: note:         def __lt__(self, int) -> bool
    

    I suggest you to play a little with this code in playground to understand this not very obvious part better.

    Now things are clear. If you define your protocol comparable only (read: at least) with same type, error will go away:

    from __future__ import annotations
    from typing import List, Protocol, TypeVar
    
    _T = TypeVar('_T')
    
    class ComparableProto(Protocol):
        def __lt__(self: _T, __other: _T) -> bool: ...  # Only same type allowed
    
    Comparable = TypeVar("Comparable", bound=ComparableProto)
    
    def sort(arr: List[Comparable]) -> List[Comparable]:
        return arr
    
    sort([1, 2, 3])
    

    Dunder name in protocols is used to explicitly say that name of argument is not important (so calling obj.__lt__(other=something) will be disallowed for protocol, but may be allowed for any subclass that has this argument named other). It's obviously the case for __lt__, because it is intended for use in operator form.

    I understand that the above example can't work, because Comparable type must be subtype of ComparableProto and just because some type implements __lt__ doesn't make it subtype of ComparableProto.

    Your understanding is wrong, every class that implements __lt__ with proper signature is subtype of ComparableProto (and "subclass" term is not so important for protocols). Protocols are used to represent structural subtyping, so another class doesn't need to inherit from protocol to be compatible with it. For instance, Sized from typing (or from collections.abc if you're modern enough and use python 3.9+) is a protocol that defines __len__(self) -> int. Any class that implements __len__, that returns int, is compatible with (read: subtype of) Sized despite not inheriting from it. See PEP544 for more formal rules (note that "subtype" term is used instead of "subclass" to avoid misunderstanding).