Home > Net >  How to hint homogenous list of instances of same subclass?
How to hint homogenous list of instances of same subclass?

Time:11-08

class A:
   pass 

class A1(A):
   pass

class A2(A):
   pass

def help_here(s: WHAT_SHOULD_I_TYPE_HERE ):
   # How to hint this function that accepts list,
   # where all elements should be instances of same subclass of A.
   # Example 1: [A1(), A1(), A1()] - good
   #            all components all elements are same class,
   #            and it is subclass of A
   # Example 2: [A2(), A2(), A2(), A2()] - valid. Same reason.
   # Example 3: [A1(), A1(), A2()] - invalid, not same classes.
   # Example 4: [1, 2] - invalid. Same class, but not subclass of A.

   if isinstance(s[0], A1):
       # If I use Union[List[A1], List[A2]] as type hint,
       # next line will have no warnings.
       # But if where will be List[A] then warning arrives:
       # Expected type 'List[A1]', got 'List[A]' instead.
       # It's Pycharm IDE.
       process_a1_list(s)


def process_a1_list(lst: List[A1]):
    pass

The only idea I have is: Union[List[A1], List[A2]].

But what if I don't know all subclasses of A?

Are there any beautiful way to hint it?

Update: Not beautiful, but let IDE (Pycharm) check types the way I want.

Update2: I need IDE type hinting solution not Runtime check. I am not sure that solution I want exists at all.

CodePudding user response:

How about this?
You make a Custom type with all the subclasses and then hint with that?

Custom = Union[List[A], List[A1], List[A2]]

def help_here(s: Custom):
    # does this work?

enter image description here

CodePudding user response:

Here is a way to dynamically build the type object you want:

from typing import List, Union


class A:
   pass 


class A1(A):
   pass


class A2(A):
   pass


def subclasses(cls: type) -> set:
    """Given a class, returns the set of all subclassess of that class (including itself)."""
    def _subclasses(cls):
        yield cls
        for subclass in cls.__subclasses__():
            yield subclass
            yield from _subclasses(subclass)
    return set(_subclasses(cls))


# Equivalent to Union[List[A], List[A1], ...], for every subclass of A
ASubclassLists = Union[tuple(List[x] for x in subclasses(A))]


def help_here(s: ASubclassLists):
   # Example 1: [A1(), A1(), A1()] - good
   #            all components all elements are same class,
   #            and it is subclass of A
   # Example 2: [A2(), A2(), A2(), A2()] - valid. Same reason.
   # Example 3: [A1(), A1(), A2()] - invalid, not same classes.
   # Example 4: [1, 2] - invalid. Same class, but not subclass of A.

   if isinstance(s[0], A1):
       # If I use Union[List[A1], List[A2]] as type hint,
       # next line will have no warnings.
       # But if where will be List[A] then warning arrives:
       # Expected type 'List[A1]', got 'List[A]' instead.
       # It's Pycharm IDE.
       process_a1_list(s)


def process_a1_list(lst: List[A1]):
    pass

For this approach you need not know the subclasses of your base class, as those can be determined automatically by Python. You need only know what your base class is.

Demonstration of the type object:

>>> Union[tuple(List[x] for x in subclasses(A))]
typing.Union[typing.List[__main__.A], typing.List[__main__.A1], typing.List[__main__.A2]]
>>> Union[List[A], List[A1], List[A2]]
typing.Union[typing.List[__main__.A], typing.List[__main__.A1], typing.List[__main__.A2]]
  • Related