CODEWITHSUNDEEP
pythonpython-typing
Paul
Paul

Reputation: 311

How to type hint a return an instance of a subclass from a list of superclass in Python

If I have a class that stores a list of superclasses, and I have a method that filters out a specific subclass types, how can I set the type hint to say that the method returns a list of the subclass?

Example:

from typing import List


class Color:
    name: str

    def __init__(self, name):
        self.name = name


class Green(Color):
    def __init__(self):
        super().__init__('green')

    def do_a_green(self):
        pass


class Blue(Color):
    def __init__(self):
        super().__init__('blue')

    def do_a_blue(self):
        pass


class ColorList:
    _color_list: List[Color]

    def get_blue_colors(self) -> List[Blue]:
        return [color for color in self._color_list
                if color.name == 'blue'] # <- error: List comprehension has
                                         # incompatible type List[Color]; expected List[Blue]

Upvotes: 0

Views: 172

Answers (1)

Robin Gugel
Robin Gugel

Reputation: 1216

The best you can do is probably with a TypeGuard.

from typing import List
from typing_extensions import TypeGuard

class Color:
    name: str

    def __init__(self, name):
        self.name = name


class Green(Color):
    def __init__(self):
        super().__init__('green')

    def do_a_green(self):
        pass


class Blue(Color):
    def __init__(self):
        super().__init__('blue')

    def do_a_blue(self):
        pass

def _is_blue(color: Color) -> TypeGuard[Blue]:
    return color.name == "blue"

class ColorList:
    _color_list: List[Color]

    def get_blue_colors(self) -> List[Blue]:
        return [color for color in self._color_list
                if _is_blue(color)]

Or use cast:

    def get_blue_colors(self) -> List[Blue]:
        return cast(List[Blue], [color for color in self._color_list
                if color.name == "blue"])

Edit:

Or use isinstance:

def get_blue_colors(self) -> List[Blue]:
        return [color for color in self._color_list
                if isinstance(color, Blue)])

Upvotes: 1

Related Questions