1010from datamodel_code_generator .__main__ import main as datamodel_codegen
1111
1212from asyncapi_python .kernel .document import Operation
13+ from asyncapi_python_codegen .parser .types import ParseContext , navigate_json_pointer
1314
1415
1516class MessageGenerator :
@@ -67,35 +68,110 @@ def _collect_message_schemas(
6768 return schemas # type: ignore[return-value]
6869
6970 def _load_component_schemas (self , spec_path : Path ) -> dict [str , Any ]:
70- """Load component schemas from the AsyncAPI specification file."""
71- try :
72- with spec_path .open ("r" ) as f :
73- spec = yaml .safe_load (f )
71+ """Load component schemas from the AsyncAPI specification file and all referenced files."""
72+ all_schemas : dict [str , Any ] = {}
73+ visited_files : set [Path ] = set ()
7474
75- components = spec . get ( "components" , {})
76- schemas = components . get ( "schemas" , {})
77- messages = components . get ( "messages" , {} )
75+ def load_schemas_from_file ( file_path : Path ) -> None :
76+ """Recursively load schemas from a file and its references."""
77+ abs_path = file_path . absolute ( )
7878
79- # Combine schemas and message payloads
80- all_schemas = {}
79+ # Avoid infinite loops
80+ if abs_path in visited_files :
81+ return
82+ visited_files .add (abs_path )
8183
82- # Add component schemas directly
83- for schema_name , schema_def in schemas . items () :
84- all_schemas [ schema_name ] = schema_def
84+ try :
85+ with abs_path . open ( "r" ) as f :
86+ spec = yaml . safe_load ( f )
8587
86- # Add message payloads from components (only if not already present from schemas)
87- for msg_name , msg_def in messages .items ():
88- if isinstance (msg_def , dict ) and "payload" in msg_def :
89- schema_name = self ._to_pascal_case (msg_name )
90- # Only add if we don't already have this schema from the schemas section
88+ components = spec .get ("components" , {})
89+ schemas = components .get ("schemas" , {})
90+ messages = components .get ("messages" , {})
91+
92+ # Add component schemas directly
93+ for schema_name , schema_def in schemas .items ():
9194 if schema_name not in all_schemas :
92- all_schemas [schema_name ] = msg_def ["payload" ]
95+ # Check if this schema is itself a reference
96+ if isinstance (schema_def , dict ) and "$ref" in schema_def :
97+ ref_value : Any = schema_def ["$ref" ] # type: ignore[misc]
98+ # Resolve the reference using ParseContext utilities
99+ if isinstance (ref_value , str ):
100+ try :
101+ context = ParseContext (abs_path )
102+ target_context = context .resolve_reference (
103+ ref_value
104+ )
105+
106+ # Load and navigate to the referenced schema
107+ with target_context .filepath .open ("r" ) as ref_file :
108+ ref_spec = yaml .safe_load (ref_file )
109+
110+ if target_context .json_pointer :
111+ resolved_schema = navigate_json_pointer (
112+ ref_spec , target_context .json_pointer
113+ )
114+ else :
115+ resolved_schema = ref_spec
116+
117+ all_schemas [schema_name ] = resolved_schema
118+ except Exception as e :
119+ print (
120+ f"Warning: Could not resolve reference { ref_value } in { abs_path } : { e } "
121+ )
122+ all_schemas [schema_name ] = schema_def
123+ else :
124+ all_schemas [schema_name ] = schema_def
125+
126+ # Add message payloads from components
127+ for msg_name , msg_def in messages .items ():
128+ if isinstance (msg_def , dict ) and "payload" in msg_def :
129+ schema_name = self ._to_pascal_case (msg_name )
130+ if schema_name not in all_schemas :
131+ all_schemas [schema_name ] = msg_def ["payload" ]
132+
133+ # Find and process all external file references
134+ self ._find_and_process_refs (
135+ spec , abs_path .parent , load_schemas_from_file
136+ )
137+
138+ except Exception as e :
139+ print (f"Warning: Could not load component schemas from { abs_path } : { e } " )
140+
141+ # Start loading from the main spec file
142+ load_schemas_from_file (spec_path )
143+
144+ return all_schemas # type: ignore[return-value]
145+
146+ def _find_and_process_refs (
147+ self , data : Any , base_dir : Path , process_file : Any
148+ ) -> None :
149+ """Recursively find all $ref entries pointing to external files."""
150+ if isinstance (data , dict ):
151+ # Check if this is a reference
152+ if "$ref" in data :
153+ ref_value : Any = data ["$ref" ] # type: ignore[misc]
154+ if isinstance (ref_value , str ) and not ref_value .startswith ("#" ):
155+ # External reference - extract file path
156+ file_part : str
157+ if "#" in ref_value :
158+ file_part = ref_value .split ("#" )[0 ]
159+ else :
160+ file_part = ref_value
161+
162+ if file_part :
163+ # Resolve relative path
164+ ref_path = (base_dir / file_part ).resolve ()
165+ process_file (ref_path )
93166
94- return all_schemas # type: ignore[return-value]
167+ # Recurse into all dict values
168+ for value in data .values (): # type: ignore[misc]
169+ self ._find_and_process_refs (value , base_dir , process_file )
95170
96- except Exception as e :
97- print (f"Warning: Could not load component schemas from { spec_path } : { e } " )
98- return {}
171+ elif isinstance (data , list ):
172+ # Recurse into all list items
173+ for item in data : # type: ignore[misc]
174+ self ._find_and_process_refs (item , base_dir , process_file )
99175
100176 def _resolve_references (self , schemas : dict [str , Any ]) -> dict [str , Any ]:
101177 """Recursively resolve $ref references to use #/$defs/... instead of #/components/schemas/..."""
@@ -105,17 +181,24 @@ def resolve_in_object(obj: Any) -> Any:
105181 resolved_obj : dict [str , Any ] = {}
106182 for key , value in obj .items (): # type: ignore[misc]
107183 if key == "$ref" and isinstance (value , str ):
108- # Transform references from #/components/schemas/... to #/$defs/...
109- if value .startswith ("#/components/schemas/" ):
110- schema_name = value .split ("/" )[- 1 ]
184+ # Extract schema name from the reference
185+ schema_name = value .split ("/" )[- 1 ]
186+
187+ # Transform all component references to #/$defs/...
188+ if "#/components/schemas/" in value :
189+ # Internal or external schema reference
111190 resolved_obj [key ] = f"#/$defs/{ schema_name } "
112- elif value . startswith ( "#/components/messages/" ) :
191+ elif "#/components/messages/" in value :
113192 # Handle message references - convert message name to PascalCase
114- msg_name = value .split ("/" )[- 1 ]
115- schema_name = self ._to_pascal_case (msg_name )
193+ schema_name = self ._to_pascal_case (schema_name )
116194 resolved_obj [key ] = f"#/$defs/{ schema_name } "
117- else :
195+ elif value .startswith ("#" ):
196+ # Other internal references, keep as-is
118197 resolved_obj [key ] = value
198+ else :
199+ # External file reference (e.g., "./commons2.yaml#/components/schemas/Foo")
200+ # Extract just the schema name and point to #/$defs
201+ resolved_obj [key ] = f"#/$defs/{ schema_name } "
119202 else :
120203 resolved_obj [key ] = resolve_in_object (value )
121204 return resolved_obj
0 commit comments