diff --git a/misc/codegen/lib/schemadefs.py b/misc/codegen/lib/schemadefs.py index ea8d6cf09ce..199043de7e6 100644 --- a/misc/codegen/lib/schemadefs.py +++ b/misc/codegen/lib/schemadefs.py @@ -1,6 +1,7 @@ from typing import ( Callable as _Callable, Dict as _Dict, + List as _List, ClassVar as _ClassVar, ) from misc.codegen.lib import schema as _schema @@ -278,7 +279,7 @@ _ = _PropertyAnnotation() drop = object() -def annotate(annotated_cls: type, replace_bases: _Dict[type, type] | None = None) -> _Callable[[type], _PropertyAnnotation]: +def annotate(annotated_cls: type, add_bases: _List[type] | None = None, replace_bases: _Dict[type, type] | None = None) -> _Callable[[type], _PropertyAnnotation]: """ Add or modify schema annotations after a class has been defined previously. @@ -295,6 +296,8 @@ def annotate(annotated_cls: type, replace_bases: _Dict[type, type] | None = None _ClassPragma(p, value=v)(annotated_cls) if replace_bases: annotated_cls.__bases__ = tuple(replace_bases.get(b, b) for b in annotated_cls.__bases__) + if add_bases: + annotated_cls.__bases__ = tuple(annotated_cls.__bases__) + tuple(add_bases) for a in dir(cls): if a.startswith(_schema.inheritable_pragma_prefix): setattr(annotated_cls, a, getattr(cls, a)) diff --git a/rust/schema/annotations.py b/rust/schema/annotations.py index 5dcd085e405..1d829951c43 100644 --- a/rust/schema/annotations.py +++ b/rust/schema/annotations.py @@ -1520,3 +1520,21 @@ class _: todo!() ``` """ + +class Callable(AstNode): + """ + A callable. Either a `Function` or a `ClosureExpr`. + """ + param_list: optional["ParamList"] | child + attrs: list["Attr"] | child + +@annotate(Function, add_bases=[Callable]) +class _: + param_list: drop + attrs: drop + + +@annotate(ClosureExpr, add_bases=[Callable]) +class _: + param_list: drop + attrs: drop