Skip to content

Commit

Permalink
Replace tomlkit with tomllib and tomli-w
Browse files Browse the repository at this point in the history
  • Loading branch information
@kislyuk
kislyuk committed Apr 13, 2024
1 parent 4e5b6ab commit b58aa2a
Show file tree
Hide file tree
Showing 4 changed files with 29 additions and 29 deletions.
6 changes: 3 additions & 3 deletions README.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,9 +118,9 @@ XML parsing vulnerabilities.

TOML support
------------
``yq`` supports `TOML <https://toml.io/>`_ as well. The ``yq`` package installs an executable, ``tomlq``, which uses the
`tomlkit library <https://github.com/sdispater/tomlkit>`_ to transcode TOML to JSON, then pipes it to ``jq``. Roundtrip
transcoding is available with the ``tomlq --toml-output``/``tomlq -t`` option.
``yq`` supports `TOML <https://toml.io/>`_ as well. The ``yq`` package installs an executable, ``tomlq``, which uses
`tomllib <https://docs.python.org/3/library/tomllib.html>`_ to transcode TOML to JSON, then pipes it to ``jq``.
Roundtrip transcoding is available with the ``tomlq --toml-output``/``tomlq -t`` option.

.. admonition:: Compatibility note

Expand Down
4 changes: 0 additions & 4 deletions pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,6 @@ per-file-ignores = {"yq/__init__.py" = ["F401"]}
module = "argcomplete.*"
ignore_missing_imports = true

[[tool.mypy.overrides]]
module = "tomlkit.*"
ignore_missing_imports = true

[[tool.mypy.overrides]]
module = "yq.version"
ignore_missing_imports = true
3 changes: 2 additions & 1 deletion setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@
install_requires=[
"PyYAML >= 5.3.1",
"xmltodict >= 0.11.0",
"tomlkit >= 0.11.6",
"tomli-w >= 1.0.0",
'tomli >= 2.0.1; python_version < "3.11"',
"argcomplete >= 1.8.1",
],
extras_require={
Expand Down
45 changes: 24 additions & 21 deletions yq/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,16 @@ def decode_docs(jq_output, json_decoder):
jq_output = jq_output[pos + 1 :]
yield doc

def get_toml_loader():
if sys.version_info >= (3, 11):
import tomllib as _toml_loader
else:
import tomli as _toml_loader
return _toml_loader

def get_xmltodict_parser():
import xmltodict
return xmltodict

def xq_cli():
cli(input_format="xml", program_name="xq")
Expand Down Expand Up @@ -221,23 +231,20 @@ def yq(
prog=program_name,
)
elif input_format == "xml":
import xmltodict

if xml_item_depth != 0:
raise Exception("xml_item_depth is not supported with xq -x")

doc = xmltodict.parse(
xml_doc = get_xmltodict_parser().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_loader = get_toml_loader()
toml_doc = toml_loader.loads(input_stream.read())
json.dump(toml_doc, json_buffer, cls=JSONDateTimeEncoder)
json_buffer.write("\n")
else:
raise Exception("Unknown input format")
Expand All @@ -260,8 +267,6 @@ def yq(
explicit_end=explicit_end,
)
elif output_format == "xml":
import xmltodict

for doc in decode_docs(jq_out, json_decoder):
if xml_root:
doc = {xml_root: doc} # type: ignore
Expand All @@ -273,7 +278,7 @@ def yq(
exit_func(msg.format(program_name))
full_document = True if xml_dtd else False
try:
xmltodict.unparse(
get_xmltodict_parser().unparse(
doc, output=output_stream, full_document=full_document, pretty=True, indent=" "
)
except ValueError as e:
Expand All @@ -283,13 +288,13 @@ def yq(
raise
output_stream.write("\n")
elif output_format == "toml":
import tomlkit
import tomli_w

for doc in decode_docs(jq_out, json_decoder):
if not isinstance(doc, dict):
msg = "{}: Error converting JSON to TOML: cannot represent non-object types at top level."
exit_func(msg.format(program_name))
tomlkit.dump(doc, output_stream)
output_stream.write(tomli_w.dumps(doc))
else:
raise Exception("Unknown output format")
else:
Expand All @@ -308,28 +313,26 @@ def yq(
prog=program_name,
)
elif input_format == "xml":
import xmltodict

def emit_entry(path, entry):
json.dump(entry, jq.stdin) # type: ignore
jq.stdin.write("\n") # type: ignore
return True

for input_stream in input_streams:
doc = xmltodict.parse(
xml_doc = get_xmltodict_parser().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
toml_doc = toml_loader.loads(input_stream.read())
json.dump(toml_doc, jq.stdin, cls=JSONDateTimeEncoder) # type: ignore
jq.stdin.write("\n") # type: ignore
else:
raise Exception("Unknown input format")
Expand Down

0 comments on commit b58aa2a

Please sign in to comment.