Subscriptable Type in Python typing

Viewed 618

I would like to typecheck that the parameter of a function is subscriptable. How do I do this with Python's typing module?

I have searched the documetation but didn't find anything. But maybe it is possible to create a custom Type. How would I do that?

2 Answers

To hint the standard __getitem__ behaviour, use the generic versions of collections.abc, such as typing.Sequence, typing.MutableSequence, typing.Mapping, or typing.MutableMapping.

from typing import Mapping

def get(container: Mapping, key):
    return container[key]

get({1: 'one', 2: 'two'}, 2)

To type hint any type that supports __getitem__, define a custom typing.Protocol with the desired behaviour.

from typing import Protocol, Any

class Lookup(Protocol):
      def __getitem__(self, key) -> Any: ...

def get(container: Lookup, key):
    return container[key]

get(['zero', 'one', 'two'], 2)

Note that the sequence and mapping types are generic, and a protocol can be define as generic as well.

from typing import Protocol, TypeVar

K = TypeVar('K', contravariant=True)
V = TypeVar('V', covariant=True)


class Lookup(Protocol[K, V]):
    def __getitem__(self, key: K) -> V: ...


def get(container: Lookup[K, V], key: K) -> V:
    return container[key]


get({1: 'one', 2: 'two'}, 2)    # succeeds type checking
get({1: 'one', 2: 'two'}, '2')  # fails type checking

Something along these lines (completely untested, though) should do it:

from typing import Protocol, TypeVar

K = TypeVar("K")
V = TypeVar("V")

class Subscriptable(Protocol[K, V]):
    def __getitem__(self, k: K) -> V:
        ...
Related