Upgrade from toml to tomlkit library

dvc/utils/serialize/_toml.py

@@ -1,9 +1,12 @@
+import logging
 from contextlib import contextmanager
 
 from funcy import reraise
 
 from ._common import ParseError, _dump_data, _load_data, _modify_data
 
+logger = logging.getLogger(__name__)
+
 
 class TOMLFileCorruptedError(ParseError):
     def __init__(self, path):
@@ -14,30 +17,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_for_update(text, path):
-    """Parses text into Python structure.
+def parse_toml(text, path, preserve_comments=False):
+    rval = _parse_toml(text, path)
+
+    if preserve_comments:
+        return rval
 
-    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
+    return rval.unwrap()
 
-    decoder = TomlPreserveCommentDecoder()
-    return parse_toml(text, path, decoder=decoder)
+
+def parse_toml_for_update(text, path):
+    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):

setup.cfg

@@ -43,6 +43,7 @@ install_requires =
     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

tests/func/params/test_show.py

@@ -44,6 +44,23 @@ def test_show_toml(tmp_dir, dvc):
     }
 
 
+def test_preserve_comments(tmp_dir, dvc):
+    from dvc.utils.serialize import _toml
+
+    contents = "# A Title [foo]\nbar = 42# meaning of life\nbaz = [1, 2]\n"
+    tmp_dir.gen("params_commented.toml", "")
+    path = (tmp_dir / "params_commented.toml").fs_path
+
+    parsed = _toml.parse_toml_for_update(contents, path)
+    with open(path, "w", encoding="utf-8") as fobj:
+        _toml._dump(parsed, fobj)
+
+    with open(path, "r", encoding="utf-8") as fobj:
+        new_contents = fobj.read()
+
+    assert new_contents == contents
+
+
 def test_show_py(tmp_dir, dvc):
     tmp_dir.gen(
         "params.py",

tests/unit/dependency/test_params.py

@@ -9,6 +9,30 @@
 DEFAULT_PARAMS_FILE = ParamsDependency.DEFAULT_PARAMS_FILE
 
 
+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)
+
+
 def test_loads_params(dvc):
     stage = Stage(dvc)
     deps = loads_params(