markdown_document.py

from dataclasses import dataclass
from functools import wraps
from typing import Callable, Concatenate, ParamSpec, TypeVar


def h1(text: str) -> str:
    """Creates an h1."""
    return "# " + text


TSuper = TypeVar("TSuper", bound="MarkdownDocument")
TRetType = TypeVar("TRetType")
Params = ParamSpec("Params")


def attach_token_function(
    func: Callable[Params, str],
) -> Callable[Concatenate[TSuper, Params], TSuper]:
    @wraps(func)
    def inner(
        self: TSuper,
        *args: Params.args,
        **kwargs: Params.kwargs,
    ) -> TSuper:
        # type hints for `self`
        self.content += func(*args, **kwargs)

        return self

    return inner


@dataclass
class MarkdownDocument:
    content: str = ""

    h1 = attach_token_function(h1)


# docstring and type hints for `.h1` are propagated
doc1 = MarkdownDocument().h1("")
#  ^? has type `MarkdownDocument`


def h2(text: str) -> str:
    """Creates an h2."""
    return "## " + text


class SubMarkdownDocument(MarkdownDocument):
    h2 = attach_token_function(h2)


# docstring and type hints for `.h1` are propagated
doc2 = SubMarkdownDocument().h1("")
#  ^? has type `SubMarkdownDocument`

# docstring and type hints for `.h1` are propagated
doc2 = doc2.h2("")
#  ^? has type `SubMarkdownDocument`