How can I auto-complete python code if i use dynamic type

Question

Here is an example below:

class Service:

    def __init__(self):
        self.data = dict()

    def get(self, clazz):
        # clazz is a class Type, and always return an instance of clazz
        return self.data[clazz]

    def register(self, clazz, *args, **kwargs):
        self.data[clazz] = clazz(*args, **kwargs)

class A:
    def method_of_a(self):
        pass

class B:
    pass

service = Service()
service.register(A)
service.register(B)

now I complete the code of service, and continue coding like

x = service.get(A)
x.method_of_a()

the code works, but there is a problem, when i was coding, the IDE(pycharm in this case) will not show the function of x when you input

x.

even i modify my code like this, it does'nt work

def get(self, clazz):
    result: clazz = self.data[clazz]
    return result

I want to know if there any way to implement Service, that IDE can recognise the type of the return value?

Answer 1

Use typing.TypeVar and typing.Type to type-hint your .get() function to say "if called with the type of T, this returns an instance of T":

from typing import TypeVar, Type

T = TypeVar('T')


class Service:

    def __init__(self):
        self.data = dict()

    def get(self, clazz: Type[T]) -> T:
        return self.data[clazz]

    def register(self, clazz, *args, **kwargs):
        self.data[clazz] = clazz(*args, **kwargs)


class A:
    def method_of_a(self):
        pass


s = Service()
s.register(A)

x = s.get(A)
x.method_of_a()
reveal_type(x)

When run with mypy, the reveal_type() call prints out

so73566368.py:24: note: Revealed type is "so73566368.A"