When using type hints combined with mypy, I want to define a generic class that allows access to an attribute that is only available for some of the generic classes. I am wondering how to define this in Python type hints.
For instance, let's have the following example:
class SomeObject:
pass
class SomeObjectWithAttr:
def __init__(self, attr):
self.attr = attr
class AppendListWithAttrs:
def __init__(self):
self._list = []
self._attrs = []
def append(self, obj, record_attr=False):
self._list.append(obj)
if record_attr:
self._attrs.append(obj.attr)
l = AppendListWithAttrs()
l.append(SomeObject())
l.append(SomeObjectWithAttr())
l.append(SomeObjectWithAttr(), record_attr=True)
I have attempted annotating this, but mypy can't seem to get happy. This is my best attempt:
from typing import Generic, TypeVar, Protocol, overload, Literal, Union
T = TypeVar("T")
class HasAttr(Protocol):
attr: str
class AppendListWithAttrs(Generic[T]):
def __init__(self):
self._list: list[T] = []
self._attrs: list[str] = []
@overload
def append(self: 'AppendListWithAttrs[HasAttr]', obj: HasAttr, record_attr: Literal[True]): ...
@overload
def append(self, obj: T, record_attr: Literal[False]): ...
def append(self, obj: T, record_attr: bool = False):
self._list.append(obj)
if record_attr:
self._attrs.append(obj.attr)
This is also not quite correct, but I'm effectively looking for a way to tell mypy that T can be anything, but if it is bounded by HasAttr, it may access this attr. Perhaps something like subclassing T with HasAttr, but maybe I'm looking at this wrong.
How would you annotate this?
How about this? It passes MyPy.
from typing import overload, Literal, Union, Any, Protocol, Generic, TypeVar
class GenericObjectProto(Protocol):
pass
class ObjectWithAttrProto(GenericObjectProto, Protocol):
attr: str
T = TypeVar('T', bound=GenericObjectProto)
class AppendListWithAttrs(Generic[T]):
def __init__(self) -> None:
self._list: list[T] = []
self._attrs: list[str] = []
@overload
def append(self, obj: GenericObjectProto, record_attr: Literal[False] = False) -> None: ...
@overload
def append(self, obj: ObjectWithAttrProto, record_attr: Literal[True] = ...) -> None: ...
def append(self, obj, record_attr: bool = False) -> None:
self._list.append(obj)
if record_attr:
self._attrs.append(obj.attr)
class SomeObject:
pass
class SomeObjectWithAttr:
def __init__(self, attr: str) -> None:
self.attr = attr
mylist = AppendListWithAttrs[SomeObject]()
mylist.append(SomeObject())
mylist.append(SomeObjectWithAttr('hello'))
mylist.append(SomeObjectWithAttr('hello'), record_attr=True)
reveal_type(mylist._list) # Revealed type is "builtins.list[__main__.SomeObject*]"
How it works
We can get some special type-hinting magic by having our ObjectWithAttr protocol inherit from our GenericObject protocol. By doing this, we have defined ObjectWithAttr as a structural subtype of GenericObject. This means that MyPy does not complain when we append an instance of SomeObjectWithAttr to a list of SomeObject instances. Because of the "structural inheritance" scheme we've defined, MyPy understands SomeObjectWithAttr to be fully compatible with a list of SomeObject instances, even though there is no actual inheritance going on.
Caveats
This passes MyPy, but only if you don't use it with the --strict setting, since I didn't annotate obj in the concrete implementation of .append. MyPy appears to be able to infer the types perfectly with the existing annotations, but it will nonetheless complain that you've left out an annotation if you run the above code as-is on the --strict setting. Even with the overloads, it doesn't like you annotating obj with T or a Union type (which seems a little buggy to me), so if you're running MyPy on --strict, you can either annotate obj in the concrete implementation with Any, or put a # type: ignore[no-untyped-def] at the end of the line.