I need to make a tree of very simple nodes (has a value and children) and I’ve found sticking to collections.abc types drastically reduces the likelihood I introduce unnecessary coupling.
@dataclasses.dataclass
class Node[TKey, TValue](MutableMapping[TKey, Self]): # Error
value: TValue
children: dict[TKey, Self]
... # abstract methods omitted for brevity
The problem is Pyright (and mypy) say I can’t use Self in the inheritance.
- Pyright: ““Self” is not valid in this context”
- mypy: “Self type is only allowed in annotations within class definition [misc]”
If python/typing#548: Higher-Kinded TypeVars were implemented in full then I’d be able to define my own ‘self’ type in the Type Parameter Scope. Which functions effectively the same way as Self, except I wouldn’t have the type auto set to self in def foo(self: TSelf) -> TSelf.
However, with current tooling we can get pretty close to a ‘perfect’ implementation with Pyright. The only issue is we can’t TSelf = Node[TKey, TValue, TSelf] in the type parameter, but we can do the equivalent in a type alias.
Should we be able to inherit with Self?
I’ve found lying to the type checker is the best stop-gap but copying a bunch of typehints from typeshed seems very suspect.
Here’s an MVCE of the different approaches:
from __future__ import annotations
from collections.abc import Iterator, MutableMapping
import dataclasses
from typing import TYPE_CHECKING, Any, Self, overload, reveal_type
"""
Don't make Node a MutableMapping.
"""
@dataclasses.dataclass
class NodeA[TKey]:
children: dict[TKey, Self]
root_a = NodeA[str]({"foo": NodeA({"bar": NodeA({"baz": NodeA({})})})})
reveal_type(root_a.children["foo"].children["bar"].children.get("baz"))
# "NodeA[str] | None"
"""
Don't allow subclassing.
"""
@dataclasses.dataclass
class NodeB[TKey](MutableMapping[TKey, "NodeB[TKey]"]):
children: dict[TKey, NodeB[TKey]]
def __getitem__(self, key: TKey) -> NodeB[TKey]:
return self.children[key]
def __setitem__(self, key: TKey, value: NodeB[TKey]) -> None:
self.children[key] = value
def __delitem__(self, key: TKey) -> None:
del self.children[key]
def __iter__(self) -> Iterator[TKey]:
return iter(self.children)
def __len__(self) -> int:
return len(self.children)
root_b = NodeB[str]({"foo": NodeB({"bar": NodeB({"baz": NodeB({})})})})
reveal_type(root_b["foo"]["bar"].get("baz"))
# "NodeB[str] | None"
"""
Use the syntactic vinegar form of [Self].
[Self]: https://docs.python.org/3/library/typing.html#typing.Self
"""
@dataclasses.dataclass
class NodeC[TKey, TSelf: "NodeC[Any, Any]" = "NodeC[TKey, TSelf]"](MutableMapping[TKey, TSelf]):
children: dict[TKey, TSelf]
def __getitem__(self, key: TKey) -> TSelf:
return self.children[key]
def __setitem__(self, key: TKey, value: TSelf) -> None:
self.children[key] = value
def __delitem__(self, key: TKey) -> None:
del self.children[key]
def __iter__(self) -> Iterator[TKey]:
return iter(self.children)
def __len__(self) -> int:
return len(self.children)
root_c0 = NodeC[str]({"foo": NodeC({"bar": NodeC({"baz": NodeC({})})})})
reveal_type(root_c0.get("foo"))
# Pyright-> "NodeC[str, Unknown] | None"
# mypy-> "NodeC[str, NodeC[str, NodeC[Any, Any]]] | None"
reveal_type(root_c0["foo"]["bar"].get("baz"))
# Pyright-> "Unknown"
# mypy-> "NodeC[Any, Any] | None"
# In Pyright: fix the self-referential TypeVar with a type alias.
type ANodeC[T] = NodeC[T, ANodeC[T]]
root_c1 = NodeC[str, ANodeC[str]]({"foo": NodeC({"bar": NodeC({"baz": NodeC({})})})})
reveal_type(root_c1["foo"]["bar"].get("baz"))
# Pyright-> "ANodeC[str] | None"
# mypy-> "Any"
"""
Use the syntactic sugar form of Self.
Given the introduction of [Type Parameter Scopes][tps] the context is well defined and is limited to the class.
We can kinda use the vinegar form without issue too, so I'm not really sure what the problem is.
[tps]: https://peps.python.org/pep-0695/#type-parameter-scopes
"""
@dataclasses.dataclass
class NodeD[TKey](MutableMapping[TKey, Self]):
children: dict[TKey, Self]
def __getitem__(self, key: TKey) -> Self:
return self.children[key]
def __setitem__(self, key: TKey, value: Self) -> None:
self.children[key] = value
def __delitem__(self, key: TKey) -> None:
del self.children[key]
def __iter__(self) -> Iterator[TKey]:
return iter(self.children)
def __len__(self) -> int:
return len(self.children)
root_d = NodeD[str]({"foo": NodeD({"bar": NodeD({"baz": NodeD({})})})})
reveal_type(root_d.get("foo"))
# Pyright-> "Unknown | None"
# mypy-> "Any | None"
reveal_type(root_d["foo"]["bar"].get("baz"))
# Pyright-> "Unknown | None"
# mypy-> "Any | None"
"""
Just lie.
"""
@dataclasses.dataclass
class NodeE[TKey](MutableMapping[TKey, Any]):
children: dict[TKey, Self]
def __getitem__(self, key: TKey) -> Self:
return self.children[key]
def __setitem__(self, key: TKey, value: Self) -> None:
self.children[key] = value
def __delitem__(self, key: TKey) -> None:
del self.children[key]
def __iter__(self) -> Iterator[TKey]:
return iter(self.children)
def __len__(self) -> int:
return len(self.children)
if TYPE_CHECKING:
@overload
def get(self, key: TKey, /) -> Self | None: ...
@overload
def get[T](self, key: TKey, /, default: Self | T) -> Self | T: ...
root_e = NodeE[str]({"foo": NodeE({"bar": NodeE({"baz": NodeE({})})})})
reveal_type(root_e["foo"]["bar"].get("baz"))
# "NodeE[str] | None"