Rust: Add Callable as a base class of Function and ClosureExpr

misc/codegen/lib/schemadefs.py

@@ -1,6 +1,7 @@
 from typing import (
     Callable as _Callable,
     Dict as _Dict,
+    Iterable as _Iterable,
     ClassVar as _ClassVar,
 )
 from misc.codegen.lib import schema as _schema
@@ -278,7 +279,7 @@ def __or__(self, other: _schema.PropertyModifier):
 drop = object()
 
 
-def annotate(annotated_cls: type, replace_bases: _Dict[type, type] | None = None) -> _Callable[[type], _PropertyAnnotation]:
+def annotate(annotated_cls: type, add_bases: _Iterable[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 decorator(cls: type) -> _PropertyAnnotation:
             _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(add_bases)
         for a in dir(cls):
             if a.startswith(_schema.inheritable_pragma_prefix):
                 setattr(annotated_cls, a, getattr(cls, a))

misc/codegen/test/test_schemaloader.py

@@ -914,6 +914,36 @@ class _:
     }
 
 
+def test_annotate_add_bases():
+    @load
+    class data:
+        class Root:
+            pass
+
+        class A(Root):
+            pass
+
+        class B(Root):
+            pass
+
+        class C(Root):
+            pass
+
+        class Derived(A):
+            pass
+
+        @defs.annotate(Derived, add_bases=(B, C))
+        class _:
+            pass
+    assert data.classes == {
+        "Root": schema.Class("Root", derived={"A", "B", "C"}),
+        "A": schema.Class("A", bases=["Root"], derived={"Derived"}),
+        "B": schema.Class("B", bases=["Root"], derived={"Derived"}),
+        "C": schema.Class("C", bases=["Root"], derived={"Derived"}),
+        "Derived": schema.Class("Derived", bases=["A", "B", "C"]),
+    }
+
+
 def test_annotate_drop_field():
     @load
     class data: