python dict where key and value depends the same generic type

user_7624196 0
Viewed 90

I want to define a dict whose key and value are share the same generic type and has some constraint with it. The following is the example of such a situation. However, applying mypy to the following code cause an error:

tmp.py:8: error: Type variable "tmp.AT" is unbound
tmp.py:8: note: (Hint: Use "Generic[AT]" or "Protocol[AT]" base class to bind "AT" inside a class)
tmp.py:8: note: (Hint: Use "AT" in function signature to bind "AT" inside a function)
Found 1 error in 1 file (checked 1 source file)

How to fix this? The reason why I need such a dict is that I want to add a type constraint between dict's key and value.

tmp.py:

from typing import Dict, Generic, TypeVar, Type, List

class A: pass
class B(A): pass
class C(A): pass

AT = TypeVar("AT", bound=A)
d: Dict[Type[AT], List[AT]] = {}
mypy typing python-3.x dictionary
1 Answers

I finally circumvent this problem by defining custom dict that inherited from the Dict.

from typing import Dict, TypeVar, Type, List
from dataclasses import dataclass

AT = TypeVar("AT")

class ConstrainedDict(Dict):
    def __getitem__(self, k: Type[AT]) -> List[AT]:
        return super().__getitem__(k)

class Foo: pass
class Bar: pass

d: ConstrainedDict
a: List[Foo] = d[Foo]  # expected to pass
b: List[Foo] = d[Bar]  # expected to cause error

mypy result is as I expected:

tmp.py:15: error: Invalid index type "Type[Bar]" for "ConstrainedDict"; expected type "Type[Foo]"
Found 1 error in 1 file (checked 1 source file)
user_76241960
Related