lib.{data,wiring}: add Python 3.14 annotation support.

Author: whitequark

amaranth/lib/data.py

@@ -2,6 +2,10 @@
 from collections.abc import Mapping, Sequence
 import warnings
 import operator
+try:
+    import annotationlib # py3.14+
+except ImportError:
+    annotationlib = None # py3.13-
 
 from amaranth._utils import final
 from amaranth.hdl import *
@@ -1208,22 +1212,36 @@ def __repr__(self):
 
 class _AggregateMeta(ShapeCastable, type):
     def __new__(metacls, name, bases, namespace):
-        if "__annotations__" not in namespace:
+        if "__annotations__" not in namespace and "__annotate_func__" not in namespace:
             # This is a base class without its own layout. It is not shape-castable, and cannot
             # be instantiated. It can be used to share behavior.
             return type.__new__(metacls, name, bases, namespace)
         elif all(not hasattr(base, "_AggregateMeta__layout") for base in bases):
+            annotations = None
+            skipped_annotations = set()
+            wrapped_annotate = None
+            if annotationlib is not None:
+                if annotate := annotationlib.get_annotate_from_class_namespace(namespace):
+                    annotations = annotationlib.call_annotate_function(
+                        annotate, format=annotationlib.Format.VALUE)
+                    def wrapped_annotate(format):
+                        annos = annotationlib.call_annotate_function(annotate, format, owner=cls)
+                        return {k: v for k, v in annos.items() if k not in skipped_annotations}
+            if annotations is None:
+                annotations = namespace.get("__annotations__", {})
+
             # This is a leaf class with its own layout. It is shape-castable and can
             # be instantiated. It can also be subclassed, and used to share layout and behavior.
             layout  = dict()
             default = dict()
-            for field_name in {**namespace["__annotations__"]}:
+            for field_name in {**annotations}:
                 try:
-                    Shape.cast(namespace["__annotations__"][field_name])
+                    Shape.cast(annotations[field_name])
                 except TypeError:
                     # Not a shape-castable annotation; leave as-is.
                     continue
-                layout[field_name] = namespace["__annotations__"].pop(field_name)
+                skipped_annotations.add(field_name)
+                layout[field_name] = annotations.pop(field_name)
                 if field_name in namespace:
                     default[field_name] = namespace.pop(field_name)
             cls = type.__new__(metacls, name, bases, namespace)
@@ -1234,6 +1252,8 @@ def __new__(metacls, name, bases, namespace):
                                      .format(", ".join(default.keys())))
             cls.__layout  = cls.__layout_cls(layout)
             cls.__default = default
+            if wrapped_annotate is not None:
+                cls.__annotate__ = wrapped_annotate
             return cls
         else:
             # This is a class that has a base class with a layout and annotations. Such a class

amaranth/lib/wiring.py

@@ -2,6 +2,10 @@
 import enum
 import re
 import warnings
+try:
+    import annotationlib # py3.14+
+except ImportError:
+    annotationlib = None # py3.13-
 
 from .. import tracer
 from ..hdl._ast import Shape, ShapeCastable, Const, Signal, Value
@@ -1669,7 +1673,14 @@ def __init__(self, signature=None, *, src_loc_at=0):
         cls = type(self)
         members = {}
         for base in reversed(cls.mro()[:cls.mro().index(Component)]):
-            for name, annot in base.__dict__.get("__annotations__", {}).items():
+            annotations = None
+            if annotationlib is not None:
+                if annotate := annotationlib.get_annotate_from_class_namespace(base.__dict__):
+                    annotations = annotationlib.call_annotate_function(
+                        annotate, format=annotationlib.Format.VALUE)
+            if annotations is None:
+                annotations = base.__dict__.get("__annotations__", {})
+            for name, annot in annotations.items():
                 if name.startswith("_"):
                     continue
                 if type(annot) is Member:

tests/test_lib_data.py

@@ -1,6 +1,11 @@
 from enum import Enum
+import sys
 import operator
 from unittest import TestCase
+try:
+    import annotationlib # py3.14+
+except ImportError:
+    annotationlib = None # py3.13-
 
 from amaranth.hdl import *
 from amaranth.lib import data
@@ -1319,7 +1324,11 @@ class S(data.Struct):
             c: str = "x"
 
         self.assertEqual(data.Layout.cast(S), data.StructLayout({"a": unsigned(1)}))
-        self.assertEqual(S.__annotations__, {"b": int, "c": str})
+        if annotationlib is not None:
+            annotations = annotationlib.get_annotations(S, format=annotationlib.Format.VALUE)
+        else:
+            annotations = S.__annotations__
+        self.assertEqual(annotations, {"b": int, "c": str})
         self.assertEqual(S.c, "x")
 
     def test_signal_like(self):