diff --git a/dvc/utils/serialize/__init__.py b/dvc/utils/serialize/__init__.py index 1dd1e92778..8b0b5490b4 100644 --- a/dvc/utils/serialize/__init__.py +++ b/dvc/utils/serialize/__init__.py @@ -1,11 +1,12 @@ from collections import defaultdict from ._common import * # noqa, pylint: disable=wildcard-import +from ._json import * # noqa, pylint: disable=wildcard-import from ._toml import * # noqa, pylint: disable=wildcard-import from ._yaml import * # noqa, pylint: disable=wildcard-import LOADERS = defaultdict(lambda: load_yaml) # noqa: F405 -LOADERS.update({".toml": load_toml}) # noqa: F405 +LOADERS.update({".toml": load_toml, ".json": load_json}) # noqa: F405 MODIFIERS = defaultdict(lambda: modify_yaml) # noqa: F405 -MODIFIERS.update({".toml": modify_toml}) # noqa: F405 +MODIFIERS.update({".toml": modify_toml, ".json": modify_json}) # noqa: F405 diff --git a/dvc/utils/serialize/_json.py b/dvc/utils/serialize/_json.py new file mode 100644 index 0000000000..3c247e85af --- /dev/null +++ b/dvc/utils/serialize/_json.py @@ -0,0 +1,29 @@ +import json + +from funcy import contextmanager, reraise + +from ._common import ParseError, _dump_data, _load_data, _modify_data + + +class JSONFileCorruptedError(ParseError): + def __init__(self, path): + super().__init__(path, "JSON file structure is corrupted") + + +def load_json(path, tree=None): + return _load_data(path, parser=parse_json, tree=tree) + + +def parse_json(text, path): + with reraise(json.JSONDecodeError, JSONFileCorruptedError(path)): + return json.loads(text) or {} + + +def dump_json(path, data, tree=None): + return _dump_data(path, data, dumper=json.dump, tree=tree) + + +@contextmanager +def modify_json(path, tree=None): + with _modify_data(path, parse_json, dump_json, tree=tree) as d: + yield d