2020from typing import (
2121 IO ,
2222 Any ,
23- TextIO ,
2423 Union ,
2524 cast ,
2625)
@@ -146,7 +145,8 @@ def my_represent_none(
146145 if output_format == "json" :
147146 json_dump (entry , output_file )
148147 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 )
150150
151151
152152def rewrite (
@@ -247,15 +247,16 @@ def rewrite_schemadef(
247247 rewrite_import (entry )
248248 elif "name" in entry and "/" in entry ["name" ]:
249249 entry_file , entry ["name" ] = entry ["name" ].lstrip ("#" ).split ("/" )
250- for field in entry [ "fields" ] :
250+ for field in entry . get ( "fields" , []) :
251251 field ["name" ] = field ["name" ].split ("/" )[2 ]
252252 rewrite_types (field , entry_file , True )
253253 with (output_dir / entry_file ).open ("a" , encoding = "utf-8" ) as entry_handle :
254254 yaml_dump (entry , entry_handle , pretty )
255255 entry ["$import" ] = entry_file
256256 del entry ["name" ]
257257 del entry ["type" ]
258- del entry ["fields" ]
258+ if "fields" in entry :
259+ del entry ["fields" ]
259260 seen_imports = set ()
260261
261262 def seen_import (entry : MutableMapping [str , Any ]) -> bool :
@@ -280,32 +281,18 @@ def json_dump(entry: Any, output_file: Path) -> None:
280281
281282def yaml_dump (
282283 entry : Any ,
283- output_file_or_handle : Union [ str , Path , TextIOWrapper , TextIO ] ,
284+ output_handle : TextIOWrapper ,
284285 pretty : bool ,
285286) -> None :
286287 """Output object as YAML."""
288+ if pretty :
289+ output_handle .write (stringify_dict (entry ))
290+ return
287291 yaml = YAML (typ = "rt" , pure = True )
288292 yaml .default_flow_style = False
289293 yaml .indent = 4
290294 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 )
309296
310297
311298if __name__ == "__main__" :
0 commit comments