20
20
from typing import (
21
21
IO ,
22
22
Any ,
23
- TextIO ,
24
23
Union ,
25
24
cast ,
26
25
)
@@ -146,7 +145,8 @@ def my_represent_none(
146
145
if output_format == "json" :
147
146
json_dump (entry , output_file )
148
147
elif output_format == "yaml" :
149
- yaml_dump (entry , output_file , pretty )
148
+ with output_file .open ("w" , encoding = "utf-8" ) as output_handle :
149
+ yaml_dump (entry , output_handle , pretty )
150
150
151
151
152
152
def rewrite (
@@ -247,15 +247,16 @@ def rewrite_schemadef(
247
247
rewrite_import (entry )
248
248
elif "name" in entry and "/" in entry ["name" ]:
249
249
entry_file , entry ["name" ] = entry ["name" ].lstrip ("#" ).split ("/" )
250
- for field in entry [ "fields" ] :
250
+ for field in entry . get ( "fields" , []) :
251
251
field ["name" ] = field ["name" ].split ("/" )[2 ]
252
252
rewrite_types (field , entry_file , True )
253
253
with (output_dir / entry_file ).open ("a" , encoding = "utf-8" ) as entry_handle :
254
254
yaml_dump (entry , entry_handle , pretty )
255
255
entry ["$import" ] = entry_file
256
256
del entry ["name" ]
257
257
del entry ["type" ]
258
- del entry ["fields" ]
258
+ if "fields" in entry :
259
+ del entry ["fields" ]
259
260
seen_imports = set ()
260
261
261
262
def seen_import (entry : MutableMapping [str , Any ]) -> bool :
@@ -280,32 +281,18 @@ def json_dump(entry: Any, output_file: Path) -> None:
280
281
281
282
def yaml_dump (
282
283
entry : Any ,
283
- output_file_or_handle : Union [ str , Path , TextIOWrapper , TextIO ] ,
284
+ output_handle : TextIOWrapper ,
284
285
pretty : bool ,
285
286
) -> None :
286
287
"""Output object as YAML."""
288
+ if pretty :
289
+ output_handle .write (stringify_dict (entry ))
290
+ return
287
291
yaml = YAML (typ = "rt" , pure = True )
288
292
yaml .default_flow_style = False
289
293
yaml .indent = 4
290
294
yaml .block_seq_indent = 2
291
-
292
- if isinstance (output_file_or_handle , (str , Path )):
293
- with open (output_file_or_handle , "w" , encoding = "utf-8" ) as result_handle :
294
- if pretty :
295
- result_handle .write (stringify_dict (entry ))
296
- return
297
- yaml .dump (entry , result_handle )
298
- return
299
- elif isinstance (output_file_or_handle , (TextIOWrapper , TextIO )):
300
- if pretty :
301
- output_file_or_handle .write (stringify_dict (entry ))
302
- return
303
- yaml .dump (entry , output_file_or_handle )
304
- return
305
- else :
306
- raise ValueError (
307
- f"output_file_or_handle must be a string or a file handle but got { type (output_file_or_handle )} "
308
- )
295
+ yaml .dump (entry , output_handle )
309
296
310
297
311
298
if __name__ == "__main__" :
0 commit comments