[Python] Generic typing for Python 3.13 CST

I'm currently writing a library that generates CST hierarchy relations for TreeSitter grammars. Each node should be typed to constrain the permitted parents and children, and common tree functionality is defined on the base node type. Generated classes then inherit from the base node types and should be extendable with language / node-specific functionality. The library is then used e.g. for writing a linter.

This is working mostly OK at runtime, but I'm having trouble with the static typing because recursive generics make my head hurt. I've boiled it down to a simplified version that exhibits the same problem. It's all a flat file, I'll just split it for comments between sections.

I'm open to any kind of refactor that makes it easier to work with, but hopefully there's some way to just alter it lightly. I'm not a dev primarily and I don't have formal education, so it's very possible that I just don't know some fundamental concept or pattern that would help here.

Base node


from __future__ import annotations

from abc import ABC
from dataclasses import dataclass
from typing import TYPE_CHECKING, ClassVar, Protocol, cast

from pydantic import BaseModel, PrivateAttr
from pydantic_core import core_schema

if TYPE_CHECKING:
from collections.abc import Iterator


class AnyNode:
def __get_pydantic_core_schema__(self, _source_type):
_ = _source_type
return core_schema.any_schema()


class NoNode:
def __get_pydantic_core_schema__(self, _source_type):
_ = _source_type
return core_schema.any_schema()


@dataclass
class WalkResult:
parent: Unknown | None
child: Unknown


Abstract = NoNode | AnyNode


class Base[P: Abstract, C: Abstract](BaseModel, AnyNode, ABC):
_id: ClassVar[str]
_typ: ClassVar[str]
_name: str
_parent: P = PrivateAttr()
_children: list[C] = PrivateAttr(default_factory=list)
_registry: ClassVar[dict[str, dict[str, type[Unknown]]]] = {}

def __init__(self, parent: P, name: str, *args, **kwargs):
super().__init__(*args, **kwargs)
self._parent = parent
self._name = name

@property
def parent(self) -> P | None:
"""Get the parent of the node."""
return self._parent

@property
def children(self) -> list[C] | None:
"""Get the children of the node."""
return self._children

def add_child(self, child: C) -> None:
self._children.append(child)

def walk(self) -> Iterator[WalkResult]:
parent = cast('Unknown', self.parent)
yield WalkResult(parent=parent, child=self)
for child in self.children or []:
child = cast('Unknown', child)
yield from child.walk()

@classmethod
def get_id(cls) -> str:
"""Get the identifier of the node class."""
return cls._id

@classmethod
def get_typ(cls) -> str:
"""Get the identifier of the node class."""
return cls._typ

@classmethod
def register(cls, sub: type[Unknown]):
sub_typ = sub.get_typ()
sub_id = sub.get_id()
if sub_typ not in cls._registry:
cls._registry[sub_typ] = {}
if sub_id in cls._registry[sub_typ]:
_err = f'Node type "{sub_id}" already registered for "{sub_typ}".'
raise ValueError(_err)
cls._registry[sub_typ][sub_id] = sub

@classmethod
def get_cls(cls, typ: str, id_: str) -> type[Unknown]:
return cls._registry[typ][id_]

def __repr__(self) -> str:
return f'{self.__class__.__name__}(name="{self._name}", typ="{self._typ}", id="{self._id}")'


def register_node(cls):
if not cls.__name__.startswith('_'):
Base.register(cls)
return cls

Unknown = Base[Abstract, Abstract]


walk() is the only tree method I have on the demo version, but of course there are many more in practice. Most things here relate to the node registry, which allow me to get the correct class after identifying language and node type.

Tree nodes


Next, there are more specific subclasses for root/branch/leaf nodes:

class Orphan[C: Abstract](Base[NoNode, C]):
_parent: NoNode = PrivateAttr(default_factory=NoNode)

@property
def parent(self) -> None:
"""Get the parent of the orphan node."""
return None

@property
def children(self) -> list[C]:
"""Get the children of the orphan node."""
return self._children


class Branch[P: Abstract, C: Abstract](Base[P, C]):
@property
def parent(self) -> P:
"""Get the parent of the branch node."""
return self._parent

@property
def children(self) -> list[C]:
"""Get the children of the branch node."""
return self._children


class Leaf[P: Abstract](Base[P, NoNode]):
_children: list[NoNode] = PrivateAttr(default_factory=list)

def add_child(self, child: NoNode) -> None:
_ = child
_err = 'Leaf nodes cannot have children.'
raise NotImplementedError(_err)

@property
def parent(self) -> P:
"""Get the parent of the leaf node."""
return self._parent

@property
def children(self) -> None:
"""Get the children of the leaf node."""
return None


These pretty much just narrow the types that each node can have.

Interfaces


Interfaces are defined as well, although I'm underusing them due to the same typing worries:

class ITree(Protocol):
def walk(self) -> Iterator[WalkResult]:
"""Walk the tree and yield nodes."""
...


class INode[P: Abstract, C: Abstract](ITree, Protocol):
@property
def parent(self) -> P | None:
"""Get the parent of the node."""
...

@property
def children(self) -> list[C] | None:
"""Get the children of the node."""
...

def add_child(self, child: C) -> None:
"""Add a child to the node."""
...


class INodeOrphan[C: Abstract](INode[NoNode, C], Protocol):
@property
def parent(self) -> None:
"""Get the parent of the orphan node."""
...

@property
def children(self) -> list[C]:
"""Get the children of the orphan node."""
...


class INodeBranch[P: Abstract, C: Abstract](INode[P, C], Protocol):
@property
def parent(self) -> P:
"""Get the parent of the branch node."""
...

@property
def children(self) -> list[C]:
"""Get the children of the branch node."""
...


class INodeLeaf[P: Abstract](INode[P, NoNode], Protocol):
@property
def parent(self) -> P:
"""Get the parent of the leaf node."""
...

@property
def children(self) -> None:
"""Get the children of the leaf node."""
return None


if TYPE_CHECKING:
_a: type[INode[Abstract, Abstract]] = Base
_b: type[INodeOrphan[Abstract]] = Orphan
_c: type[INodeBranch[Abstract, Abstract]] = Branch
_d: type[INodeLeaf[Abstract]] = Leaf


They validate correctly. ITree is split off because there are some special nodes like file/module that are only trees but not nodes - I think adding these here would only confuse things.

Base mixins


Next, there are some base mixins - they have more functionality in practice but here, they're only there to replicate the hierarchy. "Sub" would be the language-specific type in everything that follows.

class BaseMixin(BaseModel, ABC):
_mix_typ: ClassVar[str]


class _Generated__SubMixin:
_mix_typ: ClassVar[str] = 'sub_mixin'


class SubMixin(BaseMixin):
def do(self, value: str) -> None:
"""Example method in a mixin."""
print(f'This is a mixin with value: {value}')


class SubBase[P: Abstract, C: Abstract](Base[P, C], SubMixin):
_typ: ClassVar[str] = 'sub'


SubUnknown = SubBase[Abstract, Abstract]

Language tree


For each language, the base types are inherited, and the generics then define permitted parents and children. So here, a body is the root node, which can contain blocks, and each block can contain strings, which are leaves.


class SubOrphan[C: Abstract](Orphan[C], SubMixin):
_typ: ClassVar[str] = 'sub'


class SubBranch[P: Abstract, C: Abstract](Branch[P, C], SubMixin):
_typ: ClassVar[str] = 'sub'


class SubLeaf[P: Abstract](Leaf[P], SubMixin):
_typ: ClassVar[str] = 'sub'


class _Generated__SubBody(SubOrphan['SubBlock']):
_id: ClassVar[str] = 'body'


class _Generated__SubBlock(SubBranch['SubBody', 'SubString']):
_id: ClassVar[str] = 'block'


class _Generated__SubString(SubLeaf['SubBlock']):
_id: ClassVar[str] = 'string'

Node mixins - the problem


class SubBlockMixin(SubMixin):
def foo(self) -> SubUnknown:
"""Example method in a mixin."""
me = cast('SubBlock', self) if TYPE_CHECKING else self
me.do('Hello from SubBlockMixin')
print('This is a block mixin')
return next(me.walk()).child


The me cast is the only way I found to make the class behave as a SubBlock (i.e. it has the BaseNode methods / properties and the do from the SubMixin - see below for composition.) The self cannot be typed to SubBlock because it's outside the hierarchy.

Note that the foo() is also defined to return a SubUnknown. This is of course a problem, because

Diagnostics:
1. Type "Unknown" is not assignable to return type "SubUnknown"
"Base[Abstract, Abstract]" is not assignable to "SubBase[Abstract, Abstract]" [reportReturnType]


next(walk()).child is an Unknown, not a SubUnknown. Of course I can also cast here, but this is becoming really ugly. But for the outside user (the linter), it's important to have the language-specific functionality on the return value, not just a generic BaseNode.

Composition


In practice, there are a lot of these - importlib determines whether a node has a mixin and then adds it to the hierarchy. These are the classes that are supposed to be used outside the library, and they combine the functionality of the correct BaseNode subtype, the language-specific mixin and the node mixin.

@register_node
class SubBody(_Generated__SubMixin, _Generated__SubBody):
pass


@register_node
class SubBlock(SubBlockMixin, _Generated__SubMixin, _Generated__SubBlock):
pass


@register_node
class SubString(_Generated__SubMixin, _Generated__SubString):
pass

Usage


Just a minimal usage example:

def main():
sub_body = Base.get_cls('sub', 'body')(parent=NoNode(), name='root')

sub_block_a = Base.get_cls('sub', 'block')(parent=sub_body, name='block_a')
sub_block_b = Base.get_cls('sub', 'block')(parent=sub_body, name='block_b')

sub_body.add_child(sub_block_a)
sub_body.add_child(sub_block_b)

sub_string_aa = Base.get_cls('sub', 'string')(parent=sub_block_a, name='string_aa')
sub_string_ab = Base.get_cls('sub', 'string')(parent=sub_block_a, name='string_ab')
sub_string_ba = Base.get_cls('sub', 'string')(parent=sub_block_b, name='string_ba')
sub_string_bb = Base.get_cls('sub', 'string')(parent=sub_block_b, name='string_bb')

sub_block_a.add_child(sub_string_aa)
sub_block_a.add_child(sub_string_ab)
sub_block_b.add_child(sub_string_ba)
sub_block_b.add_child(sub_string_bb)

for walk_result in sub_body.walk():
print(walk_result)

print('---')

sub_string_aa = cast('SubString', sub_string_aa)
reveal_type(sub_string_aa.parent) # Type of sub_string_aa.parent is SubBlock

sub_string_aa.parent.do('Hello, World!')
ret = sub_string_aa.parent.foo()
print(f'Returned: {ret!r}')


if __name__ == '__main__':
main()


This works as expected, and the do() and foo() access are accepted by the static type checker.

WalkResult(parent=None, child=SubBody(name="root", typ="sub", id="body"))
WalkResult(parent=SubBody(name="root", typ="sub", id="body"), child=SubBlock(name="block_a", typ="sub", id="block"))
WalkResult(parent=SubBlock(name="block_a", typ="sub", id="block"), child=SubString(name="string_aa", typ="sub", id="string"))
WalkResult(parent=SubBlock(name="block_a", typ="sub", id="block"), child=SubString(name="string_ab", typ="sub", id="string"))
WalkResult(parent=SubBody(name="root", typ="sub", id="body"), child=SubBlock(name="block_b", typ="sub", id="block"))
WalkResult(parent=SubBlock(name="block_b", typ="sub", id="block"), child=SubString(name="string_ba", typ="sub", id="string"))
WalkResult(parent=SubBlock(name="block_b", typ="sub", id="block"), child=SubString(name="string_bb", typ="sub", id="string"))
---
Runtime type is 'SubBlock'
This is a mixin with value: Hello, World!
This is a mixin with value: Hello from SubBlockMixin
This is a block mixin
Returned: SubBlock(name="block_a", typ="sub", id="block")


Can anyone suggest a shed some light on how to proceed? Is it even possible to refactor this to achieve my goal, or did I bite off more than can be chewed?

I'm not tied to the existing hierarchy - in the end, I just want to

• Define as much functionality as possible on the base nodes
• Allow specifying additional functionality for languages and nodes manually
• Keep the generated code light because it's messy enough as it is
• Keep the outside interface as simple and human-compatible as possible

Thanks in advance for any advice!

Continue reading...