Type annotation for spacy custom extension

[jkgenser]

Hello, I have a simple custom extension that creates an interval tree of ents like so:

from intervaltree import Interval, IntervalTree
from spacy.tokens import Doc


def build_ent_interval_tree(doc: Doc) -> IntervalTree:
    tree = IntervalTree()
    for ent in doc.ents:
        tree.add(Interval(ent.start_char, ent.end_char))
    return tree

I set this as a method on doc like so:

Doc.set_extension("ent_interval_tree", method=build_ent_interval_tree)

What I would like to do is have type inference via my language server and enforce typing via mypy when it's being called, e.g.

ent_tree = doc._.ent_interval_tree()

I couldn't find good documentation or guidance on how to do this. Any advice/thoughts?

So below would be syntax highlighted instead of white and so the ent_tree is inferred as type IntervalTree

[NixBiks]

Did you find a good solution. Currently my workaround is to make a DocExtension class which wraps around Doc. That extension then have a method which is typed. Not a good solution but at least I get my "typesafety" using the DocExtension instead of