How to type ``NodeNG.parent``

[DanielNoord]

from __future__ import annotations

from typing import Generic

from typing_extensions import TypeVar

_NodeParent = TypeVar(
    "_NodeParent", bound="LocalsDictNode | None", default="LocalsDictNode"
)


class Node(Generic[_NodeParent]):
    def __init__(self, parent: _NodeParent) -> None:
        self.parent = parent


class LocalsDictNode(Node["LocalsDictNode" | None]):
    ...


class Module(LocalsDictNode):
    def __init__(self) -> None:
        super().__init__(None)
        self.parent: None = None
        reveal_type(self.parent)

    def method(self) -> None:
        reveal_type(self.parent)


class FunctionDef(LocalsDictNode):
    def method(self) -> None:
        if isinstance(self.parent, Module):
            reveal_type(self.parent.parent)
        else:
            reveal_type(self.parent.parent)


class Arguments(Node[FunctionDef]):
    def __init__(self, parent: FunctionDef) -> None:
        super().__init__(parent)
        reveal_type(self.parent)


class AssignName(Node[LocalsDictNode]):
    def __init__(self, parent: LocalsDictNode) -> None:
        super().__init__(parent)
        reveal_type(self.parent)

    def method(self, attr: Node) -> None:
        reveal_type(attr)

The above code almost meets all requirements we would have for .parent except for one. I'll list them in the hopes of getting any good ideas on how to fix this. Note that this uses the proposed PEP 696 as I saw no other way to even get this far without.

1. Node itself should not need any type parameters (because of its genericness) and should be useable without any
2. Unless otherwise specified NodeNG.parent should be a LocalsDictNode
3. A LocalsDictNode.parent should also be a LocalsDictNode
4. Except for Module.parent. That should be None.

The above design almost meets that requirement except for that in FunctionDef.method the second reveal_typeshowsLocalsDictNode | None, which should be LocalsDictNodeas we know thatnode.parentisn't aModuleand thereforenode.parent.parentshould beLocalsDictNode`.

/CC @cdce8p as you might have a good idea for this. Hopefully...

[cdce8p]

You could try making LocalsDictNode generic as well. The issues then just become recursive loops that crash even pyright / pylance 😄

Something like this. Note that I needed to remove the bound and default arguments as well as use Any for FunctionDef. All in all, I'm not sure this is practical tbh.

from __future__ import annotations

from typing import Generic, Any

from typing_extensions import TypeVar

_Parent = TypeVar(
    "_Parent"  #, bound="LocalsDictNode | None", default="LocalsDictNode"
)


class Node(Generic[_Parent]):
    def __init__(self, parent: _Parent) -> None:
        self.parent = parent


class LocalsDictNode(Node[_Parent]):
    ...


class Module(LocalsDictNode[None]):
    def __init__(self) -> None:
        super().__init__(None)
        self.parent = None
        reveal_type(self.parent)

    def method(self) -> None:
        reveal_type(self.parent)


class FunctionDef(LocalsDictNode[Node[LocalsDictNode[Any | None]]]):
    def method(self) -> None:
        if isinstance(self.parent, Module):
            reveal_type(self.parent.parent)  # None
        else:
            reveal_type(self.parent.parent)  # LocalsDictNode[Any | None]

[cdce8p]

Another suggestion would be to use a Protocol type instead of LocalsDictNode for typing. That could help prevent circular imports.

[DanielNoord]

Yeah I thought about the recursiveness, but that didn't really feel good as well..

I wonder if Protocols could be our saviour. Have you ever tried making a NodeNGProtocol and using the everywhere?