CODEWITHSUNDEEP
pythonpython-typingmypy
PBauma
PBauma

Reputation: 13

Python define return type based on property type

I have three classes, one base class and two child classes with the following structure:

class BaseLayout:
    def add_element(self: "BaseLayout", element: str) -> None:
        ...

class LayoutA(BaseLayout):
    def add_element(self: "BaseLayout", element: str, name: str) -> None:
        ...

class LayoutB(BaseLayout):
    def add_element(self: "BaseLayout", element: str, number: int) -> None:
        ...

The base class is not directly used but does implement some other functions which are not relevant to the question. I then use the layout as property in the Container class:

class Container:
    def __init__(self: "Container", layout: BaseLayout) -> None:
        self._current_layout = layout

    @property
    def layout(self: "Container") -> BaseLayout:
        return self._current_layout

Here is my current problem:

If I provided an instance of LayoutA to the Container class and then call the add_element method I get a type checker error Expected 1 positional argument. Here is the code to add the element:

c = Container(LayoutA())

# Expected 1 positional argument
c.layout.add_element("foo", "bar")

How can I provided the correct type hints to the Container class to make this work?

Upvotes: 0

Views: 316

Answers (1)

Daniil Fajnberg
Daniil Fajnberg

Reputation: 18388

As explained already by @chepner in the comments, your code is first of all not type safe because your BaseLayout subclasses' add_element overrides are not compatible with the base signature.

You could fix that for example by creating different methods, if you need different behavior, and have them call the superclass' method. I don't know the actual use case, so this is just an example. There are many ways to deal with this.

class BaseLayout:
    def add_element(self, element: str) -> None:
        print(f"{element=}")


class LayoutA(BaseLayout):
    def add_name_element(self, element: str, name: str) -> None:
        super().add_element(element)
        print(f"{name=}")


class LayoutB(BaseLayout):
    def add_number_element(self, element: str, number: int) -> None:
        super().add_element(element)
        print(f"{number=}")

Once that is fixed, you can make Container generic in terms of the layout it uses like this:

from typing import Generic, TypeVar


L = TypeVar("L", bound=BaseLayout)


class Container(Generic[L]):
    _current_layout: L  # this is optional, but helpful IMO

    def __init__(self, layout: L) -> None:
        self._current_layout = layout

    @property
    def layout(self) -> L:
        return self._current_layout


ca = Container(LayoutA())
ca.layout.add_name_element("foo", "bar")

cb = Container(LayoutB())
cb.layout.add_number_element("foo", 1)

This is fine and static type checkers should correctly infer the layout type. If you add reveal_type(ca._current_layout) and reveal_type(cb._current_layout) below and run mypy, you'll get this:

note: Revealed type is "LayoutA"
note: Revealed type is "LayoutB"

This also means if you e.g. try and call cb.layout.add_name_element somewhere, the type checker will complain.

Upvotes: 1

Related Questions