为可迭代类添加类型注解：解决MyElement迭代元素类型识别问题

你的代码运行正常但PyCharm报错，是因为ET.Element的迭代器默认类型注解为Iterator[ET.Element]，静态类型检查器无法识别实际迭代返回的是MyElement实例。要解决这个问题，只需在MyElement类中重写__iter__方法，明确指定返回的迭代器元素类型为MyElement。

解决方案代码 #

from __future__ import annotations
import xml.etree.ElementTree as ET
from typing import Iterator, cast


class MyElement(ET.Element):
    def my_method(self):
        print('OK')
    
    # 重写__iter__方法，指定迭代器返回MyElement类型
    def __iter__(self) -> Iterator[MyElement]:
        return super().__iter__()


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

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

root.my_method()

for ch in root:
    ch.my_method()  # PyCharm不再报错

补充说明 #

如果需要支持通过索引访问子元素（比如root[0]）也能正确识别类型，可以额外重写__getitem__方法：

from typing import Union, slice

class MyElement(ET.Element):
    # ... 其他方法 ...
    
    def __getitem__(self, index: Union[int, slice]) -> Union[MyElement, list[MyElement]]:
        return super().__getitem__(index)

这样无论是迭代还是索引访问，PyCharm都能正确识别子元素的类型为MyElement。

内容的提问来源于stack exchange，提问作者gog