Annotating return types for methods returning self in mixins

Question

I am using a builder pattern where most methods on a (big) class return their identity (self) and are thus annotated to return the type of the class they're a member of:

class TextBuilder:
    parts: List[str]           # omitted
    render: Callable[[], str]  # for brevity

    def text(self, val: str) -> "TextBuilder":
        self.parts.append(val)
        return self

    def bold(self, val: str) -> "TextBuilder":
        self.parts.append(f"<b>{val}</b>")
        return self

    ...

Example usage:

joined_text = TextBuilder().text("a ").bold("bold").text(" text").render()  
# a <b>bold</b> text

Now as this class is growing large I would like to split and group related methods up into mixins:

class BaseBuilder:
    parts: List[str]           # omitted
    render: Callable[[], str]  # for brevity


class TextBuilder(BaseBuilder):
    def text(self, val: str):
        self.parts.append(val)
        return self
    ...


class HtmlBuilder(BaseBuilder):
    def bold(self, val: str):
        self.parts.append(f"<b>{val}</b>")
        return self
    ...


class FinalBuilder(TextBuilder, HtmlBuilder):
    pass

However, I do not see a way to properly annotate the mixin classes' return types in a way that the resulting class FinalBuilder always makes mypy believe that it returns FinalBuilder and not one of the mixin classes. All that of course assuming I want to actually annotate self and return types because they may not be inferred from what goes on inside those methods.

I have tried making the mixin classes generic and marking them explicitly as returning a type T bound to BaseBuilder, but that did not satisfy mypy. Any ideas? For now I am just going to skip all these shenanigans and omit the return types everywhere as they should be properly inferred when using the FinalBuilder, but I'm still curious if there is a general way to approach this.

Answer 1

If you want the return type to always be what self is, just annotate the self parameter like so:

from typing import List, Callable, TypeVar

T = TypeVar('T', bound=BaseBuilder)

class BaseBuilder:
    parts: List[str]           # omitted
    render: Callable[[], str]  # for brevity


class TextBuilder(BaseBuilder):
    def text(self: T, val: str) -> T:
        self.parts.append(val)
        return self
    ...


class HtmlBuilder(BaseBuilder):
    def bold(self: T, val: str) -> T:
        self.parts.append(f"<b>{val}</b>")
        return self
    ...


class FinalBuilder(TextBuilder, HtmlBuilder):
    pass


# Type checks
f = FinalBuilder().text("foo").bold("bar")

# Mypy states this is type 'FinalBuilder'
reveal_type(f)

A few notes:

1. If we don't annotate self, mypy will normally assume it's the type of whatever class we're currently contained in. However, it's actually fine to give it a custom type hint if you want, so long as that type hint is compatible with the class. (For example, it wouldn't be legal to add a def foo(self: int) -> None to HtmlBuilder since int isn't a supertype of HtmlBuilder.)

   We take advantage of this by making self generic so we can specify a more specific return type.

   See the mypy docs for more details: https://mypy.readthedocs.io/en/stable/generics.html#generic-methods-and-generic-self

2. I bounded the TypeVar to BaseBuilder so that both functions would be able to see the parts and render fields. If you want your text(...) and bold(...) functions to also see fields defined within TextBuilder and HtmlBuilder respectively, you'll need to create two TypeVars bound to these more specific child classes.