diff --git a/dvc/utils/serialize/_toml.py b/dvc/utils/serialize/_toml.py index 11fcaf832..31313896a 100644 --- a/dvc/utils/serialize/_toml.py +++ b/dvc/utils/serialize/_toml.py @@ -14,30 +14,31 @@ def load_toml(path, fs=None): return _load_data(path, parser=parse_toml, fs=fs) -def parse_toml(text, path, decoder=None): - from toml import TomlDecodeError, loads +def _parse_toml(text, path): + from tomlkit import loads + from tomlkit.exceptions import ParseError as TomlkitParseError - with reraise(TomlDecodeError, TOMLFileCorruptedError(path)): - return loads(text, decoder=decoder) + with reraise(TomlkitParseError, TOMLFileCorruptedError(path)): + return loads(text) + + +def parse_toml(text, path, preserve_comments=False): + rval = _parse_toml(text, path) + + if preserve_comments: + return rval + + return rval.unwrap() def parse_toml_for_update(text, path): - """Parses text into Python structure. - - NOTE: Python toml package does not currently use ordered dicts, so - keys may be re-ordered between load/dump, but this function will at - least preserve comments. - """ - from toml import TomlPreserveCommentDecoder - - decoder = TomlPreserveCommentDecoder() - return parse_toml(text, path, decoder=decoder) + return parse_toml(text, path, preserve_comments=True) -def _dump(data, stream): - import toml +def _dump(data, stream, sort_keys=False): + import tomlkit - return toml.dump(data, stream, encoder=toml.TomlPreserveCommentEncoder()) + return tomlkit.dump(data, stream, sort_keys=sort_keys) def dump_toml(path, data, fs=None, **kwargs): diff --git a/setup.cfg b/setup.cfg index c30ff7489..bf7b7f5d2 100644 --- a/setup.cfg +++ b/setup.cfg @@ -42,7 +42,7 @@ install_requires = distro>=1.3.0 appdirs>=1.4.3 ruamel.yaml>=0.17.11 - toml>=0.10.1 + tomlkit>=0.11.1 funcy>=1.14 pathspec>=0.9.0,<0.10.0 shortuuid>=0.5.0 diff --git a/tests/unit/utils/serialize/test_toml.py b/tests/unit/utils/serialize/test_toml.py new file mode 100644 index 000000000..b6e04ffb3 --- /dev/null +++ b/tests/unit/utils/serialize/test_toml.py @@ -0,0 +1,38 @@ +def test_preserve_comments(tmp_dir): + from dvc.utils.serialize._toml import modify_toml + + contents_fmt = """\ +#A Title +[foo] +bar = {} # meaning of life +baz = [1, 2] +""" + tmp_dir.gen("params.toml", contents_fmt.format("42")) + + with modify_toml("params.toml") as d: + d["foo"]["bar"] //= 2 + assert (tmp_dir / "params.toml").read_text() == contents_fmt.format("21") + + +def test_parse_toml_type(): + from tomlkit.toml_document import TOMLDocument + + from dvc.utils.serialize._toml import parse_toml + + contents = "# A Title [foo]\nbar = 42# meaning of life\nbaz = [1, 2]\n" + + parsed = parse_toml(contents, ".") + assert not isinstance(parsed, TOMLDocument) + assert isinstance(parsed, dict) + + +def test_parse_toml_for_update(): + from tomlkit.toml_document import TOMLDocument + + from dvc.utils.serialize._toml import parse_toml_for_update + + contents = "# A Title [foo]\nbar = 42# meaning of life\nbaz = [1, 2]\n" + + parsed = parse_toml_for_update(contents, ".") + assert isinstance(parsed, TOMLDocument) + assert isinstance(parsed, dict)