Add type hints: mwparserfromhell.nodes.heading

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)