1+
2+ from collections .abc import Generator
13from types import MethodType
2- import pydantic
3- from typing import TypeVar
4+ from typing import (
5+ Tuple , TypeVar ,
6+ )
47from typing_extensions import is_typeddict
5- _T_TypedDict = TypeVar ('_T_TypedDict' )
68
7- def amaranth_annotate (modeltype : type ['_T_TypedDict' ], schema_id : str , member = '__chipflow_annotation__' , decorate_object = False ):
9+ import pydantic
10+ from amaranth import Fragment
11+ from amaranth .lib import meta , wiring
12+
13+
14+ _T_TypedDict = TypeVar ('_T_TypedDict' )
15+ def amaranth_annotate (modeltype : type [_T_TypedDict ], schema_id : str , member : str = '__chipflow_annotation__' , decorate_object = False ):
16+ # a bit of nastyness as can't set TypedDict as a bound yet
817 if not is_typeddict (modeltype ):
9- raise TypeError (f''' amaranth_annotate must be passed a TypedDict, not { modeltype } ''' )
18+ raise TypeError (f" amaranth_annotate must be passed a TypedDict, not { modeltype } " )
1019
1120 # interesting pydantic issue gets hit if arbitrary_types_allowed is False
1221 if hasattr (modeltype , '__pydantic_config__' ):
13- config = getattr (modeltype , '__pydantic_config__' )
22+ config : pydantic . ConfigDict = getattr (modeltype , '__pydantic_config__' )
1423 config ['arbitrary_types_allowed' ] = True
1524 else :
1625 config = pydantic .ConfigDict ()
1726 config ['arbitrary_types_allowed' ] = True
1827 setattr (modeltype , '__pydantic_config__' , config )
28+
1929 PydanticModel = pydantic .TypeAdapter (modeltype )
2030
2131 def annotation_schema ():
2232 schema = PydanticModel .json_schema ()
23- schema ['$schema' ] = ' https://json-schema.org/draft/2020-12/schema'
33+ schema ['$schema' ] = " https://json-schema.org/draft/2020-12/schema"
2434 schema ['$id' ] = schema_id
2535 return schema
2636
27- class Annotation :
28- ' Generated annotation class'
37+ class Annotation ( meta . Annotation ) :
38+ " Generated annotation class"
2939 schema = annotation_schema ()
3040
3141 def __init__ (self , parent ):
3242 self .parent = parent
3343
34- def origin (self ):
44+ @property
45+ def origin (self ): # type: ignore
3546 return self .parent
3647
37- def as_json (self ):
38- return PydanticModel .dump_python (getattr (self .parent , member ))
48+ def as_json (self ): # type: ignore
49+ # TODO: this is slow, but atm necessary as dump_python doesn't do the appropriate
50+ # transformation of things like PosixPath. Figure out why, maybe log issue/PR with
51+ # pydantic
52+ # return json.loads(PydanticModel.dump_json(getattr(self.parent, member)))
53+ return PydanticModel .dump_python (getattr (self .parent , member ), mode = 'json' )
3954
4055 def decorate_class (klass ):
4156 if hasattr (klass , 'annotations' ):
4257 old_annotations = klass .annotations
4358 else :
4459 old_annotations = None
45-
46- def annotations (self , obj ):
60+ def annotations (self , obj , / ): # type: ignore
4761 if old_annotations :
48- annotations = old_annotations (self , obj )
62+ annotations = old_annotations (self , obj ) # type: ignore
4963 else :
5064 annotations = super (klass , obj ).annotations (obj )
5165 annotation = Annotation (self )
52- return annotations + (annotation ,)
66+ return annotations + (annotation ,) # type: ignore
67+
5368
5469 klass .annotations = annotations
5570 return klass
@@ -60,19 +75,46 @@ def decorate_obj(obj):
6075 else :
6176 old_annotations = None
6277
63- def annotations (self = None , origin = None ):
78+ def annotations (self , origin , / ): # type: ignore
6479 if old_annotations :
6580 annotations = old_annotations (origin )
6681 else :
6782 annotations = super (obj .__class__ , obj ).annotations (obj )
6883 annotation = Annotation (self )
69- return annotations + (annotation ,)
84+ return annotations + (annotation ,) # type: ignore
7085
71- setattr (obj , 'annotations' , MethodType (annotations , obj ))
86+ setattr (obj , 'annotations' , MethodType (annotations , obj ))
7287 return obj
7388
7489 if decorate_object :
7590 return decorate_obj
7691 else :
7792 return decorate_class
7893
94+
95+ def submodule_metadata (fragment : Fragment , component_name : str , recursive = False ) -> Generator [Tuple [wiring .Component , str | tuple , dict ]]:
96+ """
97+ Generator that finds `component_name` in `fragment` and
98+ then yields the ``wiring.Component``s of that component's submodule, along with their names and metadata
99+
100+ Can only be run once for a given component (or its children)
101+
102+ If recursive = True, then name is a tuple of the heirarchy of names
103+ otherwise, name is the string name of the first level component
104+ """
105+
106+ subfrag = fragment .find_subfragment (component_name )
107+ design = subfrag .prepare ()
108+ for k ,v in design .elaboratables .items ():
109+ full_name :tuple = design .fragments [design .elaboratables [k ]].name
110+ if len (full_name ) > 1 : # ignore the top component
111+ if recursive :
112+ name = full_name [1 :]
113+ else :
114+ if len (full_name ) != 2 :
115+ continue
116+ name = full_name [1 ]
117+ if isinstance (k , wiring .Component ):
118+ metadata = k .metadata .as_json ()['interface' ]
119+ yield k , name , metadata
120+
0 commit comments