How to annotate a return type as either a class instance or its (unique) subclass instance?-CodePudding

I am writing a python library used by importing and (optionally) sub-classing some of the 'helper classes' it provides. I fail to come up with a design that would properly let static analysis tools properly recognise the types that my 'helper classes' methods deal with. Here's a MWE illustrating (one of) the issues I run into:

My lib

from typing import Dict


class Thing:
    def shout(self):
        print(f"{self} says AAAAAAAAAaaaaaaaa")


class ContainerOfThings:
    def __init__(self):
        thing_cls = self._thing_cls = get_unique_subclass(Thing)
        self._things: Dict[str, thing_cls] = {}

    def add_thing(self, id_: str):
        self._things[id_] = self._thing_cls()

    def get_thing(self, id_: str):
        return self._things[id_]


def get_unique_subclass(cls):
    # this works but maybe there's a better way to do this?
    classes = cls.__subclasses__()
    if len(classes) == 0:
        return cls
    elif len(classes) == 1:
        return classes[0]
    elif len(classes) > 1:
        raise RuntimeError(
            "This class should only be subclassed once", cls, classes
        )

What I expect users to do with it

class BetterThing(Thing):
    def be_civilized(self):
        print(f"{self} says howdy!")

container = ContainerOfThings()
container.add_thing("some_id")
thingy = container.get_thing("some_id")
thingy.be_civilized()
thingy.do_something_invalid()  # here I would like mypy to detect that this will not work

This snippet does not alarm static analysis tools, because thingy is detected as Any, but fails at runtime on the last line because do_something_invalid() is not defined. Isn't it possible to give hints that thingy is in fact an instance of BetterThing here?

My attempts so far:

Attempt 1

Annotate ContainerOfThings._things as Dict[str, Thing] instead of Dict[str, thing_cls]

This passes mypy, but pycharm detects thingy as an instance of Thing and thus complains about "Unresolved attribute reference 'be_civilized' for class 'Thing'"

Attempt 2

Annotate ContainerOfThings.get_thing() return value as Thing

Less surprisingly, this triggers errors from both pycharm and mypy about Thing not having the 'be_civilized' attribute.

Attempt 3

Use ThingType = TypeVar("ThingType", bound=Thing) as return value for ContainerOfThings.get_thing()

I believe (?) that this is what TypeVar is intended for, and it works, except for the fact that mypy then requires thingy to be be annotated with BetterThing, along with every return value of ContainerOfThings.get_thing(), which will be quite cumbersome with my 'real' library.

Is there an elegant solution for this? Is get_unique_subclass() too dirty a trick to play nice with static analysis? Is there something clever to do with typing_extensions.Protocol that I could not come up with?

Thanks for your suggestions.

CodePudding user response：

Basically you need ContainerOfThings to be generic:
https://mypy.readthedocs.io/en/stable/generics.html#defining-generic-classes

And then I think it would be better for ContainerOfThings to be explicit about the type of thing that it will generate instead of auto-magically locating some sub-class that has been defined.

We can put this together in a way that will satisfy mypy (and I would expect pycharm too, though I haven't tried it)...

from typing import Dict, Generic, Type, TypeVar


class Thing:
    def shout(self):
        print(f"{self} says AAAAAAAAAaaaaaaaa")


T = TypeVar('T', bound=Thing)


class ContainerOfThings(Generic[T]):
    def __init__(self, thing_cls: Type[T]):
        self._thing_cls = thing_cls
        self._things: Dict[str, T] = {}

    def add_thing(self, id_: str):
        self._things[id_] = self._thing_cls()

    def get_thing(self, id_: str) -> T:
        return self._things[id_]


class BetterThing(Thing):
    def be_civilized(self):
        print(f"{self} says howdy!")
        

container = ContainerOfThings(BetterThing)
container.add_thing("some_id")
thingy = container.get_thing("some_id")
thingy.be_civilized()  # OK
thingy.do_something_invalid()  # error: "BetterThing" has no attribute "do_something_invalid"