diff --git a/src/mwparserfromhell/nodes/heading.py b/src/mwparserfromhell/nodes/heading.py index b65cf29..9c1b122 100644 --- a/src/mwparserfromhell/nodes/heading.py +++ b/src/mwparserfromhell/nodes/heading.py @@ -19,50 +19,58 @@ # SOFTWARE. +from typing import TYPE_CHECKING, Any, Generator, Callable, Optional + from ._base import Node from ..utils import parse_anything -__all__ = ["Heading"] +if TYPE_CHECKING: + from ..wikicode import Wikicode class Heading(Node): """Represents a section heading in wikicode, like ``== Foo ==``.""" - def __init__(self, title, level): + def __init__(self, title: Any, level: int): super().__init__() self.title = title self.level = level - def __str__(self): + def __str__(self) -> str: return ("=" * self.level) + str(self.title) + ("=" * self.level) - def __children__(self): + def __children__(self) -> Generator["Wikicode", None, None]: yield self.title - def __strip__(self, **kwargs): + def __strip__(self, **kwargs: Any) -> Optional[str]: return self.title.strip_code(**kwargs) - def __showtree__(self, write, get, mark): + def __showtree__( + self, + write: Callable[[str], None], + get: Callable[["Wikicode"], None], + mark: Callable[[], None], + ) -> None: write("=" * self.level) get(self.title) write("=" * self.level) @property - def title(self): + def title(self) -> "Wikicode": """The title of the heading, as a :class:`.Wikicode` object.""" return self._title + @title.setter + def title(self, value: Any) -> None: + self._title = parse_anything(value) + @property - def level(self): + def level(self) -> int: """The heading level, as an integer between 1 and 6, inclusive.""" return self._level - @title.setter - def title(self, value): - self._title = parse_anything(value) - @level.setter - def level(self, value): + def level(self, value: int) -> None: value = int(value) if value < 1 or value > 6: raise ValueError(value)