Merge pull request #327 from tatuylonen/overload-find-child

Overload WikiNode.find_child with a LEVEL variant
tatuylonen · Oct 15, 2024 · a709d4b · a709d4b
2 parents 706074a + a178279
commit a709d4b
Showing 1 changed file with 22 additions and 0 deletions.
diff --git a/src/wikitextprocessor/parser.py b/src/wikitextprocessor/parser.py
@@ -240,6 +240,14 @@ class NodeKind(enum.Flag):
     "======": NodeKind.LEVEL6,
 }
 
+LITERAL_LEVEL_KINDS = Literal[
+    NodeKind.LEVEL1
+    | NodeKind.LEVEL2
+    | NodeKind.LEVEL3
+    | NodeKind.LEVEL4
+    | NodeKind.LEVEL5
+    | NodeKind.LEVEL6
+]
 
 # Maps subtitle node kind to its level.  Keys include all title/subtitle nodes
 # (this is also used like a set of all subtitle kinds, including the root).
@@ -356,6 +364,20 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return self.__str__()
 
+    @overload
+    def find_child(
+        self,
+        target_kinds: LITERAL_LEVEL_KINDS,
+        with_index: Literal[True],
+    ) -> Iterator[tuple[int, "LevelNode"]]: ...
+
+    @overload
+    def find_child(
+        self,
+        target_kinds: LITERAL_LEVEL_KINDS,
+        with_index: Literal[False] = ...,
+    ) -> Iterator["LevelNode"]: ...
+
     @overload
     def find_child(
         self,