Specifying additional subclassed bounds of a TypeVar in Python typing

user_560023 0
Viewed 301

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?

python-typing mypy type-hinting structural-typing python
1 Answers

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.

user_139900160
Related