Add stuff __init__.py, typing stuff

tatuylonen · Sep 12, 2024 · f566de1 · f566de1
1 parent c451455
commit f566de1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/src/wikitextprocessor/__init__.py b/src/wikitextprocessor/__init__.py
@@ -1,12 +1,16 @@
 from .common import MAGIC_FIRST, MAGIC_LAST
-from .core import Page, Wtp
-from .parser import NodeKind, WikiNode
+from .core import Page, TemplateArgs, Wtp
+from .parser import HTMLNode, LevelNode, NodeKind, TemplateNode, WikiNode
 
 __all__ = (
     "Wtp",
+    "HTMLNode",
+    "LevelNode",
     "NodeKind",
+    "TemplateNode",
     "WikiNode",
     "MAGIC_FIRST",  # Some applications with to use the same ranges
     "MAGIC_LAST",
     "Page",
+    "TemplateArgs",
 )
diff --git a/src/wikitextprocessor/parser.py b/src/wikitextprocessor/parser.py
@@ -370,6 +370,20 @@ def find_child(
         with_index: Literal[False] = ...,
     ) -> Iterator["TemplateNode"]: ...
 
+    @overload
+    def find_child(
+        self,
+        target_kinds: Literal[NodeKind.HTML],
+        with_index: Literal[True],
+    ) -> Iterator[tuple[int, "HTMLNode"]]: ...
+
+    @overload
+    def find_child(
+        self,
+        target_kinds: Literal[NodeKind.HTML],
+        with_index: Literal[False] = ...,
+    ) -> Iterator["HTMLNode"]: ...
+
     @overload
     def find_child(
         self, target_kinds: NodeKind, with_index: Literal[True]
@@ -464,7 +478,7 @@ def find_html(
         with_index: Literal[True],
         attr_name: str,
         attr_value: str,
-    ) -> Iterator["HTMLNode"]: ...
+    ) -> Iterator[tuple[int, "HTMLNode"]]: ...
 
     @overload
     def find_html(
@@ -473,7 +487,7 @@ def find_html(
         with_index: Literal[False] = ...,
         attr_name: str = ...,
         attr_value: str = ...,
-    ) -> Iterator[tuple[int, "HTMLNode"]]: ...
+    ) -> Iterator["HTMLNode"]: ...
 
     def find_html(
         self,