Nested generics in python using the returns package

Question

As an example, I have a zoo. A zoo (Z) has a list of cages (C), and a cage has a list of animals (A).

Since mypy doesnt support nested generics properly, I am experimenting with the returns package's higher kinded types in order to get proper typechecks.

As a special type of zoo, I have an OpenZoo which can only contain OpenCages with Dogs in it.

My attempt looks like so:

import itertools

from typing import TypeVar

from returns.primitives.hkt import Kind1, Kind2, SupportsKind1, SupportsKind2

A = TypeVar('A', bound='Animal')
C = TypeVar('C', bound='Cage')
Z = TypeVar('Z', bound='Zoo')

class Animal:

    def feed(self):
        ...

class Dog(Animal):

    def bark(self):
        print("Woof!")

class Cage(SupportsKind1['Cage', A]):
    def __init__(self, animals: list[A]):
        self.animals = animals


class OpenCage(Cage[Dog]):
    def unleash(self) -> None:
        ...

class Zoo(SupportsKind2['Zoo', C, A]):

    def __init__(self, cages: list[Kind1[C, A]]):
        self.cages = cages

    def all_animals(self) -> list[A]:
        return list(itertools.chain.from_iterable([c.animals for c in iter(self.cages)]))


class OpenZoo(Zoo[OpenCage, Dog]):
    def unlock_cages(self) -> None:
        ...


my_animal = Dog()
my_cage = OpenCage([my_animal])
my_zoo = OpenZoo([my_cage]) # <--- mypy error here

However, this yields mypy error error: List item 0 has incompatible type "OpenCage"; expected "KindN[OpenCage, Dog, Any, Any]" [list-item]

Any idea how I should fix this issue? Or if there is another (better) way of achieving this?

I'm using python 3.10 and mypy 1.0.1

Answer 1

I did some research and i found this :

The error message suggests that the type of my_cage is OpenCage, which is not compatible with the expected type of KindN[OpenCage, Dog, Any, Any] for the second type parameter of OpenZoo.

This is because OpenZoo is defined as a Zoo that only contains OpenCage with Dog as the animal type, but my_cage is of type OpenCage[Dog] only. Therefore, the type of my_cage needs to be lifted to a Kind1[OpenCage, Dog] using SupportsKind1 in order to match the type parameter of Zoo.

The code below should resolve the error

import itertools
    
from typing import TypeVar

A = TypeVar('A', bound='Animal')
C = TypeVar('C', bound='Cage')
Z = TypeVar('Z', bound=Zoo[C, A])

class Animal:
    def feed(self):
        ...

class Dog(Animal):
    def bark(self):
        print("Woof!")

class Cage:
    def __init__(self, animals: list[A]):
        self.animals = animals

class OpenCage(Cage):
    def unleash(self) -> None:
        ...

class Zoo:
    def __init__(self, cages: list[Cage]):
        self.cages = cages

    def all_animals(self) -> list[A]:
        return [animal for cage in self.cages for animal in cage.animals]

class OpenZoo(Zoo):
    def unlock_cages(self) -> None:
        ...

my_animal = Dog()
my_cage = OpenCage([my_animal])
my_zoo = OpenZoo([my_cage])

Hope this helps :)