How to annotate a class method parameter with the class it self

Question

I have this error message Unresolved reference 'TreeNode' on def add_children(self, child: TreeNode) I think this an IDE (PyCharm) issue. The code run's well.

I don't understand why, and how can I fix it?

class TreeNode:
    def __init__(
        self,
        element: Entity,
        children=None
    ):
        if children is None:
            self.children = []
        else:
            self.children = children
        self.element = element

    def add_children(self, child: TreeNode):
        self.children.append(child)

score 2 · Accepted Answer · answered Aug 05 '20 at 11:31

TreeNode is only defined after the class TreeNode body. Use from __future__ import annotations or child: 'TreeNode' to make a forward reference.

class TreeNode:
    def __init__(
        self,
        element: Entity,
        children=None
    ):
        if children is None:
            self.children = []
        else:
            self.children = children
        self.element = element

    def add_children(self, child: 'TreeNode'):
        self.children.append(child)

How to annotate a class method parameter with the class it self

1 Answers1