Type annotation for an iterable class

gog 0
Viewed 33

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?

typing python
1 Answers

A couple of workarounds

1. Annotation cheese

There is a somewhat cheesy trick that doesn't violate any syntax: annotating variable before the loop.

ch: MyElement
for ch in root:
    ch.my_method()

https://peps.python.org/pep-0526/#where-annotations-aren-t-allowed

You can't annotate inside the for loop but can do so before it.

2. Override methods

You can add overridden methods for __iter__ and __next__ for your MyElement class with annotations:

def __iter__(self) -> MyElement:
    return super().iter()

def __next__(self) -> MyElement:
    pass

Element parent doesn't have __iter__ or __next__ natively so it's a bit sketchy but I didn't want to spend more time researching it so "posted as is because it worked"™ They were both required for the annotation to work. Also for prerequisites you will need to import

from __future__ import annotations

in the beginning of your script to access -> MyElement inside itself without issues.

LTJ0
Related