1
+
2
+ from collections .abc import Generator
1
3
from types import MethodType
2
- import pydantic
3
- from typing import TypeVar
4
+ from typing import (
5
+ Tuple , TypeVar ,
6
+ )
4
7
from typing_extensions import is_typeddict
5
- _T_TypedDict = TypeVar ('_T_TypedDict' )
6
8
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
8
17
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 } " )
10
19
11
20
# interesting pydantic issue gets hit if arbitrary_types_allowed is False
12
21
if hasattr (modeltype , '__pydantic_config__' ):
13
- config = getattr (modeltype , '__pydantic_config__' )
22
+ config : pydantic . ConfigDict = getattr (modeltype , '__pydantic_config__' )
14
23
config ['arbitrary_types_allowed' ] = True
15
24
else :
16
25
config = pydantic .ConfigDict ()
17
26
config ['arbitrary_types_allowed' ] = True
18
27
setattr (modeltype , '__pydantic_config__' , config )
28
+
19
29
PydanticModel = pydantic .TypeAdapter (modeltype )
20
30
21
31
def annotation_schema ():
22
32
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"
24
34
schema ['$id' ] = schema_id
25
35
return schema
26
36
27
- class Annotation :
28
- ' Generated annotation class'
37
+ class Annotation ( meta . Annotation ) :
38
+ " Generated annotation class"
29
39
schema = annotation_schema ()
30
40
31
41
def __init__ (self , parent ):
32
42
self .parent = parent
33
43
34
- def origin (self ):
44
+ @property
45
+ def origin (self ): # type: ignore
35
46
return self .parent
36
47
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' )
39
54
40
55
def decorate_class (klass ):
41
56
if hasattr (klass , 'annotations' ):
42
57
old_annotations = klass .annotations
43
58
else :
44
59
old_annotations = None
45
-
46
- def annotations (self , obj ):
60
+ def annotations (self , obj , / ): # type: ignore
47
61
if old_annotations :
48
- annotations = old_annotations (self , obj )
62
+ annotations = old_annotations (self , obj ) # type: ignore
49
63
else :
50
64
annotations = super (klass , obj ).annotations (obj )
51
65
annotation = Annotation (self )
52
- return annotations + (annotation ,)
66
+ return annotations + (annotation ,) # type: ignore
67
+
53
68
54
69
klass .annotations = annotations
55
70
return klass
@@ -60,19 +75,46 @@ def decorate_obj(obj):
60
75
else :
61
76
old_annotations = None
62
77
63
- def annotations (self = None , origin = None ):
78
+ def annotations (self , origin , / ): # type: ignore
64
79
if old_annotations :
65
80
annotations = old_annotations (origin )
66
81
else :
67
82
annotations = super (obj .__class__ , obj ).annotations (obj )
68
83
annotation = Annotation (self )
69
- return annotations + (annotation ,)
84
+ return annotations + (annotation ,) # type: ignore
70
85
71
- setattr (obj , 'annotations' , MethodType (annotations , obj ))
86
+ setattr (obj , 'annotations' , MethodType (annotations , obj ))
72
87
return obj
73
88
74
89
if decorate_object :
75
90
return decorate_obj
76
91
else :
77
92
return decorate_class
78
93
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