Search code examples
pythonpycharm

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


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)

Solution

  • 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)