From f52763a0b3e854e813840bb449bc34cacfcd4c92 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Mon, 31 Aug 2020 03:15:52 +0300
Subject: [PATCH] dvc: use flatten-dict instead of flatten-json  (#4500)

* dvc: use flatten-dict instead of flatten-json

Fixes #4499

* utils: add [un]flatten

* pyinstaller: add flatten-dict hook
---
 dvc/command/experiments.py              | 9 +++------
 dvc/command/metrics.py                  | 5 ++---
 dvc/repo/reproduce.py                   | 4 ++--
 dvc/utils/diff.py                       | 6 +++---
 dvc/utils/flatten.py                    | 9 +++++++++
 scripts/hooks/hook-dvc.utils.flatten.py | 5 +++++
 setup.py                                | 2 +-
 7 files changed, 25 insertions(+), 15 deletions(-)
 create mode 100644 dvc/utils/flatten.py
 create mode 100644 scripts/hooks/hook-dvc.utils.flatten.py

diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py
index d6343afc7..ce60c030b 100644
--- a/dvc/command/experiments.py
+++ b/dvc/command/experiments.py
@@ -9,6 +9,7 @@ from typing import Iterable, Optional
 from dvc.command.base import CmdBase, append_doc_link, fix_subparsers
 from dvc.command.metrics import DEFAULT_PRECISION
 from dvc.exceptions import DvcException, InvalidArgumentError
+from dvc.utils.flatten import flatten
 
 logger = logging.getLogger(__name__)
 
@@ -55,11 +56,9 @@ def _filter_names(
 
 
 def _update_names(names, items):
-    from flatten_json import flatten
-
     for name, item in items:
         if isinstance(item, dict):
-            item = flatten(item, ".")
+            item = flatten(item)
             names.update(item.keys())
         else:
             names.add(name)
@@ -134,8 +133,6 @@ def _format_time(timestamp):
 
 
 def _extend_row(row, names, items, precision):
-    from flatten_json import flatten
-
     def _round(val):
         if isinstance(val, float):
             return round(val, precision)
@@ -148,7 +145,7 @@ def _extend_row(row, names, items, precision):
 
     for fname, item in items:
         if isinstance(item, dict):
-            item = flatten(item, ".")
+            item = flatten(item)
         else:
             item = {fname: item}
         for name in names:
diff --git a/dvc/command/metrics.py b/dvc/command/metrics.py
index 432f17345..52eb364e5 100644
--- a/dvc/command/metrics.py
+++ b/dvc/command/metrics.py
@@ -14,9 +14,8 @@ DEFAULT_PRECISION = 5
 def show_metrics(
     metrics, all_branches=False, all_tags=False, all_commits=False
 ):
-    from flatten_json import flatten
-
     from dvc.utils.diff import format_dict
+    from dvc.utils.flatten import flatten
 
     # When `metrics` contains a `None` key, it means that some files
     # specified as `targets` in `repo.metrics.show` didn't contain any metrics.
@@ -32,7 +31,7 @@ def show_metrics(
                 continue
 
             logger.info(f"\t{fname}:")
-            for key, value in flatten(format_dict(metric), ".").items():
+            for key, value in flatten(format_dict(metric)).items():
                 logger.info(f"\t\t{key}: {value}")
 
     if missing:
diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py
index cd8bc17ef..5b3be7b51 100644
--- a/dvc/repo/reproduce.py
+++ b/dvc/repo/reproduce.py
@@ -119,10 +119,10 @@ def reproduce(
 
 
 def _parse_params(path_params):
-    from flatten_json import unflatten
     from ruamel.yaml import YAMLError
 
     from dvc.dependency.param import ParamsDependency
+    from dvc.utils.flatten import unflatten
     from dvc.utils.serialize import loads_yaml
 
     ret = {}
@@ -141,7 +141,7 @@ def _parse_params(path_params):
                 )
         if not path:
             path = ParamsDependency.DEFAULT_PARAMS_FILE
-        ret[path] = unflatten(params, ".")
+        ret[path] = unflatten(params)
     return ret
 
 
diff --git a/dvc/utils/diff.py b/dvc/utils/diff.py
index 7442564a4..219691359 100644
--- a/dvc/utils/diff.py
+++ b/dvc/utils/diff.py
@@ -1,6 +1,8 @@
 import json
 from collections import defaultdict
 
+from .flatten import flatten
+
 
 def _parse(raw):
     if raw is None or isinstance(raw, (dict, list, int, float)):
@@ -35,9 +37,7 @@ def _flatten(d):
         return defaultdict(lambda: None)
 
     if isinstance(d, dict):
-        from flatten_json import flatten as fltn
-
-        return defaultdict(lambda: None, fltn(d, "."))
+        return defaultdict(lambda: None, flatten(d))
 
     return defaultdict(lambda: "unable to parse")
 
diff --git a/dvc/utils/flatten.py b/dvc/utils/flatten.py
new file mode 100644
index 000000000..ff5f2c9f1
--- /dev/null
+++ b/dvc/utils/flatten.py
@@ -0,0 +1,9 @@
+import flatten_dict
+
+
+def flatten(d):
+    return flatten_dict.flatten(d, reducer="dot")
+
+
+def unflatten(d):
+    return flatten_dict.unflatten(d, splitter="dot")
diff --git a/scripts/hooks/hook-dvc.utils.flatten.py b/scripts/hooks/hook-dvc.utils.flatten.py
new file mode 100644
index 000000000..a4f7faf4e
--- /dev/null
+++ b/scripts/hooks/hook-dvc.utils.flatten.py
@@ -0,0 +1,5 @@
+from PyInstaller.utils.hooks import (  # pylint:disable=import-error
+    copy_metadata,
+)
+
+datas = copy_metadata("flatten-dict")
diff --git a/setup.py b/setup.py
index bdd357dd4..93366f9cb 100644
--- a/setup.py
+++ b/setup.py
@@ -75,7 +75,7 @@ install_requires = [
     "pydot>=1.2.4",
     "speedcopy>=2.0.1; python_version < '3.8' and sys_platform == 'win32'",
     "dataclasses; python_version < '3.7'",
-    "flatten_json>=0.1.6,<0.1.8",
+    "flatten_dict>=0.3.0,<1",
     "tabulate>=0.8.7",
     "pygtrie==2.3.2",
     "dpath>=2.0.1,<3",