Search code examples
pythonpycharmpython-typing

Python Annotations: return type as an instance


Does there is a way to let the return value know its type should belong to some isinstance?

for example,def myfunction(...) -> instance[MyClass]: (It's an informal grammar, I just example I want the syntax is something like that)

that means it only accepts the return type which belongs MyClass.


An example

I was given an example under the following, to make you more clear what I use for it to do what.

consider the following:

from typing import List, TypeVar, Union, Type


class Table:
    NAME: str

    def show_info(self):
        print(self.NAME)


class TableA(Table):
    NAME = 'A'

    def method_belong_a_only(self):
        print('a only')


class TableB(Table):
    NAME = 'B'

    def method_belong_b_only(self):
        print('b only')


class DBBase:
    Tables: List[Table]


class MyDB(DBBase):
    Tables = [TableA, TableB]

    def get_table(self, table_name: str) -> Union[Table, None]:
        table_list = [_ for _ in self.Tables if hasattr(_, 'NAME') and _.NAME == table_name]

        if not table_list:
            return None

        if len(table_list) > 1:
            raise ValueError('Table name conflict!')
        return table_list[0]()  # create an instance

and then if I write the following,

db = MyDB()
ta = db.get_table('A')
ta.show_info()  # ok
# ta.  # IDE will show method, which belongs to Table only.
ta.method_belong_a_only()  # It still can run, but IDE show Cannot find reference 'method_belong_a_only' in Table|None

tb: TableB = db.get_table('B')  # It still can run, but IDE show ``Expected type 'TableA', got 'Optional[Table]' instead``
# tb.  # Now, It can show ``method_belong_b_only``

enter image description here

I don't want it(PyCharm IDE) to show warnings to me, and I want it works well on intellisence.


Solution

  • You can use typing.TypeVar PEP484.type-aliases

    from typing import List, TypeVar, Union, Type
    
    
    T_Table = TypeVar('T_Table', bound=Table)
    
    
    class MyDB2(DBBase):
        Tables = [TableA, TableB]
    
        def get_table(self, t: Type[T_Table] = Table) -> Union[T_Table, None]:
            t_set = {cur_t for cur_t in self.Tables if cur_t.__name__ == t.__name__}
    
            if not t_set:
                return None
    
            table = t_set.pop()
            return table()   # create an instance
    
        def get_table_2(self, table_type_name: str) -> Union[T_Table, None]:
            t_set = {cur_t for cur_t in self.Tables if cur_t.__name__ == table_type_name}
    
            if not t_set:
                return None
    
            table = t_set.pop()
            return table()   # create an instance
    
    
    db = MyDB2()
    a = db.get_table(TableA)
    #  a = db.get_table(str)  # <-- Expected type 'Type[T_Table]', got 'Type[str]' instead
    a.method_belong_a_only()  # intellisense will work perfectly.
    
    b: TableB = db.get_table_2('TableB')  # This is another way, if IDE felt confused, you can tell him what the type it is.
    # b.  # intellisense will work perfectly.