Type annotation for an iterable class

Question

I've got a class that extends ElementTree.Element:

import xml.etree.ElementTree as ET
from typing import cast


class MyElement(ET.Element):
    def my_method(self):
        print('OK')


xml = '''<test> <sub/> <sub/> </test>'''

root: MyElement = cast(
    MyElement,
    ET.fromstring(xml, parser=ET.XMLParser(target=ET.TreeBuilder(element_factory=MyElement))))

root.my_method()  # this is fine

for ch in root:
    ch.my_method()  # PyCharm error message  ???

This does work, however the last line is highlighted by PyCharm because it considers ch to be Element, not MyElement.

How should I annotate MyElement to make it clear that when I iterate it, I get MyElement instances and not ET.Elements?

Answer 1

Annotate the __iter__ of MyElement as a method to return Iterator[MyElement]. There is almost no additional runtime overhead. Pycharm and mypy will both pass:

from collections.abc import Callable, Iterator


class MyElement(ET.Element):
    def my_method(self):
        print('OK')

    __iter__: Callable[..., Iterator['MyElement']]