Search code examples
pythonpycharmcontainerssubclasstype-hinting

How to type-hint that container subclass always contains particular type?


I'm struggling to get type-hinting to work as I would expect regarding the content types for custom container subclasses in PyCharm. Let's start with a case that works as I would expect. You can create a subclass of list and specify that it will always have int content, and Pycharm then recognizes that each item in such a list will be an int.

class IntList(list[int]): pass
il = IntList(())
answer1 = il[0]

When I mouse-over answer1 Pycharm says that it expects this to have type int, presumably because the class declaration specified that an IntList wouldn't be any old list, but instead would be a list[int]. (Never mind the fact that this code would raise an error when run because il would be empty. This was just a minimal example to show that PyCharm can sometimes draw type-hinting information from the bracketed [int] in the class declaration. The same issues arise in other cases where getting an item from the container wouldn't raise an error.)

So this worked fine when subclassing from list. But what I want to do is to create my own generic container class -- call it Box -- that could contain a variety of different types of objects. Then I want to declare my own subclass IntBox that will contain only int items, and I want PyCharm to recognize this in its various mouse-over hints, auto-completion suggestions, and linting error detection, just like it could for IntList. So here's a very pared-down example of what I'd like.

class Box(list): pass
class IntBox(Box[int]): pass
ib = IntBox(())
answer2 = ib[0]

In this case, when I mouse-over answer2, PyCharm says that it could have type Any and does not recognize that the [int] implies that this is not just a generic Box/list, but instead one whose contents have been type-hinted to be int.

I've tried all the variations I can imagine using typing.TypeVar and typing.Generic to try to more explicitly indicate that each subclass of Box will have a single type of contents, that Box.__getitem__ will return that type, and that, for the subclass IntBox that type is int.

The only solution I've found is that, when I create ib I can explicitly declare that this instance has type IntBox[int] and then PyCharm will know to expect that ib[0] will be an int. But it seems like I shouldn't need to explicitly say that each time I create an IntBox instance, and instead there should be some way to get PyCharm to infer this from the [int] in the class declaration for IntBox just like it could for IntList.

Of course this was just a toy example. In the real case that's motivating this, I want my generic container class "Box" to define other methods (not just __getitem__) that are type-hinted to return whichever specific type of object the subclass of "Box" in question always contains, where this varies across subclasses. Using TypeVar and Generic I can get this to work, if I explicitly type-declare that each subclass instance will contain a particular [contenttype], but I can't find a way to get it to work without tedious explicit type-declarations on each instance.

Edit: since solutions that work in the simple case of telling what elements will be in a list sub-subclass apparently don't automatically scale up to this real case, here's an example closer to what I need, including that Box is a nested Sequence rather than a simple list, including a Box.get_first() method that should also receive type-hinting of int for IntBox, and including what I think is roughly the right use of TypeVar:

from typing import TypeVar, Sequence
T = TypeVar('T')
class Box(Sequence[Sequence[T]]):
    def get_first(self:'Box[T]')->T:
        return self[0][0]
class IntBox(Box[int]): pass
ib = IntBox()  # works only if I declare this is type: IntBox[int]
answer = ib.get_first()  # hovering over answer should show it will be int

Further edit: the problem with the preceding seems to be the nested Sequence[Sequence[T]]. Changing this to Generic[T] makes things work as expected.


Solution

  • Use typing.Generic to pass a type to the Box superclass from a subclass using a typing.TypeVar

    from typing import Generic, TypeVar
    
    T = TypeVar('T')
    
    
    class Box(Generic[T], list[T]):
        pass
    
    
    class IntBox(Box[int]):
        pass
    
    
    class StrBox(Box[str]):
        pass
    
    
    ib = IntBox(())
    answer1 = ib[0]  # int
    sb = StrBox(())
    answer2 = sb[0]  # str