Search code examples
pythonpython-typingcovariance

How to use type-hints for a covariant mutable collection-like class in Python?


I'm trying to create a covariant collection-like class in Python, like so:

from typing import Generic, TypeVar, List, cast

class Animal():
    pass

class Dog(Animal):
    pass

class Cat(Animal):
    pass

class Zoo():
    def __init__(self, items: List[Animal]):
        self._items = items.copy()  # type: List[Animal]

    def add(self, animal: Animal) -> None:
        self._items.append(animal)

    def animal_count(self) -> int:
        return len(self._items)

    def get_animals(self) -> List[Animal]:
        return self._items.copy()

class DogHouse(Zoo):
    def __init__(self, items: List[Dog]):
        self._items = items.copy()  # type: List[Dog]

    def add(self, dog: Dog) -> None:
        assert isinstance(dog, Dog)
        self._items.append(dog)

In short, I like to subclass Zoo, so that it only takes a specific type of Animal. In this case, a DogHouse with only Dogs.

Mypy gives two error with this code:

  • error: Incompatible types in assignment (expression has type "List[Dog]", base class "Zoo" defined the type as "List[Animal]")
  • error: Argument 1 of "add" is incompatible with supertype "Zoo"; supertype defines the argument type as "Animal"

I understand what mypy is trying to warn me about: the following snippet is syntactically valid, but may lead to issue because there suddenly may be another animal (Cat, Kangaroo, ...) in the DogHouse (Formally, the code may violate the Liskov substitution principle):

doghouse = DogHouse([])
doghouse._items.append(Cat())

However, my code is supposed to take care of this, for example by checking the type in DogHouse.add(), making Zoo._items (somewhat) private, and make copious copy() of the mutable sequence, so Zoo._items can't be modified.

Is there a way to both make DogHouse a subclass of Zoo (and benefit from generic methods in Zoo), as well as to use type hinting to verify that my code won't accidental allow Cats or other Animals sneak into the DogHouse?

I've read https://mypy.readthedocs.io/en/stable/generics.html#variance-of-generics, but have trouble applying this advices to my code (coming from a duck-type language like Python, I'm not yet very verbose with the concept of covariance).


Edit: I attempted a solution by defining Animal_co = TypeVar('Animal_co', bound=Animal, covariant=True), but that lead to an error: Cannot use a covariant type variable as a parameter. See the accepted answer for the right answer, and explanation why this was wrong.


Solution

  • The attempt you made with a covariant type variable in an earlier edit was close, but the type variable shouldn't be covariant. Making it covariant would mean that a Zoo[Dog] is also a Zoo[Animal], and in particular, it would mean that add(self, animal: Animal_co) can take any Animal no matter what Animal_co is bound to. The behavior you're looking for is really invariant, not covariant. (You might want a separate "read-only" zoo ABC or Protocol that actually is covariant.)

    While you're at it, stop poking the parent's implementation details:

    T = TypeVar('T', bound=Animal)
    
    class Zoo(Generic[T]):
        _items: List[T]
        def __init__(self, items: Iterable[T]):
            self._items = list(items)
    
        def add(self, animal: T) -> None:
            self._items.append(animal)
    
        def animal_count(self) -> int:
            return len(self._items)
    
        def get_animals(self) -> List[T]:
            return self._items.copy()
    
    class DogHouse(Zoo[Dog]):
        def add(self, dog: Dog) -> None:
            assert isinstance(dog, Dog)
            super().add(dog)
    

    The assert is there for type-checking at runtime. If you don't actually care about runtime enforcement, you can simplify DogHouse to

    class DogHouse(Zoo[Dog]): pass
    

    or remove it entirely and use Zoo[Dog] directly.