Upgrade from toml to tomlkit library (#7468)

* Upgrade from toml to tomlkit library

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

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

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)