Fix return type of ElementTree.iterparse() (#14213)

Author: srittau

stdlib/xml/etree/ElementTree.pyi

@@ -249,28 +249,42 @@ def dump(elem: Element[Any] | ElementTree[Any]) -> None: ...
 def indent(tree: Element[Any] | ElementTree[Any], space: str = "  ", level: int = 0) -> None: ...
 def parse(source: _FileRead, parser: XMLParser[Any] | None = None) -> ElementTree[Element]: ...
 
-# This class is defined inside the body of iterparse
+# The type of the second element of the tuple yielded by iterparse depends
+# on the event type in the first element of the tuple:
+#  * start, end: Element[str]
+#  * comment, pi: Element[_ElementCallable]
+#  * start-ns: tuple[str, str] (prefix, uri)
+#  * end-ns: None
+_EventT_co = TypeVar("_EventT_co", bound=Element[str] | Element[_ElementCallable] | tuple[str, str] | None, covariant=True)
+_EventType: TypeAlias = Literal["start", "end", "comment", "pi", "start-ns", "end-ns"]
+
+# This class is defined inside the body of iterparse.
 @type_check_only
-class _IterParseIterator(Iterator[tuple[str, Element]], Protocol):
-    def __next__(self) -> tuple[str, Element]: ...
+class _IterParseIterator(Iterator[tuple[_EventType, _EventT_co]], Protocol[_EventT_co]):
     if sys.version_info >= (3, 13):
         def close(self) -> None: ...
     if sys.version_info >= (3, 11):
         def __del__(self) -> None: ...
 
+# See the comment for _EventT_co above for possible iterator types.
+@overload
+def iterparse(source: _FileRead, events: Iterable[_EventType]) -> _IterParseIterator[Any]: ...
 @overload
-def iterparse(source: _FileRead, events: Sequence[str] | None = None) -> _IterParseIterator: ...
+def iterparse(source: _FileRead, events: None = None) -> _IterParseIterator[Element[str]]: ...
+
+# In case a custom parser is passed, the type of the second element of the tuple
+# yielded by iterparse depends on the parser.
 @overload
 @deprecated("The `parser` parameter is deprecated since Python 3.4.")
-def iterparse(source: _FileRead, events: Sequence[str] | None = None, parser: XMLParser | None = None) -> _IterParseIterator: ...
+def iterparse(source: _FileRead, events: Iterable[_EventType], parser: XMLParser | None = None) -> _IterParseIterator[Any]: ...
 
 _EventQueue: TypeAlias = tuple[str] | tuple[str, tuple[str, str]] | tuple[str, None]
 
-class XMLPullParser(Generic[_E]):
-    def __init__(self, events: Sequence[str] | None = None, *, _parser: XMLParser[_E] | None = None) -> None: ...
+class XMLPullParser(Generic[_EventT_co]):
+    def __init__(self, events: Iterable[_EventType] | None = None, *, _parser: XMLParser[_EventT_co] | None = None) -> None: ...
     def feed(self, data: str | ReadableBuffer) -> None: ...
     def close(self) -> None: ...
-    def read_events(self) -> Iterator[_EventQueue | tuple[str, _E]]: ...
+    def read_events(self) -> Iterator[_EventQueue | tuple[_EventType, _EventT_co]]: ...
     def flush(self) -> None: ...
 
 def XML(text: str | ReadableBuffer, parser: XMLParser | None = None) -> Element: ...

stubs/defusedxml/defusedxml/ElementTree.pyi

@@ -1,6 +1,6 @@
 from _typeshed import ReadableBuffer
 from collections.abc import Sequence
-from typing import Final
+from typing import Any, Final
 from xml.etree.ElementTree import (
     Element,
     ElementTree,
@@ -49,7 +49,7 @@ XMLTreeBuilder = DefusedXMLParser
 XMLParse = DefusedXMLParser
 XMLParser = DefusedXMLParser
 
-# wrapper to xml.etree.ElementTree.parse
+# Wrapper to xml.etree.ElementTree.parse
 def parse(
     source: _FileRead,
     parser: XMLParser | None = None,
@@ -58,15 +58,17 @@ def parse(
     forbid_external: bool = True,
 ) -> ElementTree: ...
 
-# wrapper to xml.etree.ElementTree.iterparse
+# Wrapper to xml.etree.ElementTree.iterparse
+#
+# See there for possible return types.
 def iterparse(
     source: _FileRead,
     events: Sequence[str] | None = None,
     parser: XMLParser | None = None,
     forbid_dtd: bool = False,
     forbid_entities: bool = True,
     forbid_external: bool = True,
-) -> _IterParseIterator: ...
+) -> _IterParseIterator[Any]: ...
 def fromstring(
     text: str | ReadableBuffer, forbid_dtd: bool = False, forbid_entities: bool = True, forbid_external: bool = True
 ) -> Element: ...