From b284a2a7c0aa778832b1b52798b546beb2bdae6e Mon Sep 17 00:00:00 2001
From: Tom Hvitved <hvitved@github.com>
Date: Fri, 11 Oct 2024 12:00:30 +0200
Subject: [PATCH] Rust: Add `Callable` as a base class of `Function` and
 `ClosureExpr`

---
 misc/codegen/lib/schemadefs.py |  5 ++++-
 rust/schema/annotations.py     | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

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