diff --git a/yq/__init__.py b/yq/__init__.py index 2613b29..d95ac4b 100644 --- a/yq/__init__.py +++ b/yq/__init__.py @@ -42,6 +42,17 @@ def decode_docs(jq_output, json_decoder): yield doc +def get_toml_loader(): + if sys.version_info >= (3, 11): + import tomllib + + return tomllib.loads + else: + import tomlkit + + return tomlkit.parse + + def xq_cli(): cli(input_format="xml", program_name="xq") @@ -197,6 +208,8 @@ def yq( msg = "{}: Error starting jq: {}: {}. Is jq installed and available on PATH?" exit_func(msg.format(program_name, type(e).__name__, e)) + assert jq.stdin is not None # this is to keep mypy happy + try: if converting_output: # TODO: enable true streaming in this branch (with asyncio, asyncproc, a multi-shot variant of @@ -226,18 +239,18 @@ def yq( if xml_item_depth != 0: raise Exception("xml_item_depth is not supported with xq -x") - doc = xmltodict.parse( + xml_doc = xmltodict.parse( input_stream.buffer if isinstance(input_stream, io.TextIOWrapper) else input_stream.read(), disable_entities=True, force_list=xml_force_list, ) - json.dump(doc, json_buffer, cls=JSONDateTimeEncoder) + json.dump(xml_doc, json_buffer, cls=JSONDateTimeEncoder) json_buffer.write("\n") elif input_format == "toml": import tomlkit - doc = tomlkit.load(input_stream) # type: ignore - json.dump(doc, json_buffer, cls=JSONDateTimeEncoder) + toml_doc = tomlkit.load(input_stream) + json.dump(toml_doc, json_buffer, cls=JSONDateTimeEncoder) json_buffer.write("\n") else: raise Exception("Unknown input format") @@ -264,7 +277,7 @@ def yq( for doc in decode_docs(jq_out, json_decoder): if xml_root: - doc = {xml_root: doc} # type: ignore + doc = {xml_root: doc} elif not isinstance(doc, dict): msg = ( "{}: Error converting JSON to XML: cannot represent non-object types at top level. " @@ -311,31 +324,32 @@ def yq( import xmltodict def emit_entry(path, entry): - json.dump(entry, jq.stdin) # type: ignore - jq.stdin.write("\n") # type: ignore + assert jq.stdin is not None # this is to keep mypy happy + json.dump(entry, jq.stdin) + jq.stdin.write("\n") return True for input_stream in input_streams: - doc = xmltodict.parse( + xml_doc = xmltodict.parse( input_stream.buffer if isinstance(input_stream, io.TextIOWrapper) else input_stream.read(), disable_entities=True, force_list=xml_force_list, item_depth=xml_item_depth, item_callback=emit_entry, ) - if doc: - emit_entry(None, doc) + if xml_doc: + emit_entry(None, xml_doc) elif input_format == "toml": - import tomlkit - + toml_loader = get_toml_loader() for input_stream in input_streams: - json.dump(tomlkit.load(input_stream), jq.stdin, cls=JSONDateTimeEncoder) # type: ignore - jq.stdin.write("\n") # type: ignore + toml_doc = toml_loader(input_stream.read()) + json.dump(toml_doc, jq.stdin, cls=JSONDateTimeEncoder) + jq.stdin.write("\n") else: raise Exception("Unknown input format") try: - jq.stdin.close() # type: ignore + jq.stdin.close() except Exception: pass jq.wait()