Use tomllib on Python 3.11 when not round-tripping

1 files changed, 29 insertions, 15 deletions

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()