sqlite_utils.utils.flatten() function, closes #500

Author: simonw

docs/reference.rst

@@ -101,3 +101,10 @@ sqlite_utils.utils.chunks
 -------------------------
 
 .. autofunction:: sqlite_utils.utils.chunks
+
+.. _reference_utils_flatten:
+
+sqlite_utils.utils.flatten
+--------------------------
+
+.. autofunction:: sqlite_utils.utils.flatten

sqlite_utils/cli.py

@@ -24,6 +24,7 @@
     chunks,
     file_progress,
     find_spatialite,
+    flatten as _flatten,
     sqlite3,
     decode_base64_values,
     progressbar,
@@ -997,7 +998,7 @@ def insert_upsert_implementation(
                     "Invalid JSON - use --csv for CSV or --tsv for TSV files"
                 )
             if flatten:
-                docs = (dict(_flatten(doc)) for doc in docs)
+                docs = (_flatten(doc) for doc in docs)
 
     if convert:
         variable = "row"
@@ -1079,15 +1080,6 @@ def insert_upsert_implementation(
         db[table].transform(types=tracker.types)
 
 
-def _flatten(d):
-    for key, value in d.items():
-        if isinstance(value, dict):
-            for key2, value2 in _flatten(value):
-                yield key + "_" + key2, value2
-        else:
-            yield key, value
-
-
 def _find_variables(tb, vars):
     to_find = list(vars)
     found = {}
@@ -1845,7 +1837,7 @@ def memory(
             tracker = TypeTracker()
             rows = tracker.wrap(rows)
         if flatten:
-            rows = (dict(_flatten(row)) for row in rows)
+            rows = (_flatten(row) for row in rows)
         db[csv_table].insert_all(rows, alter=True)
         if tracker is not None:
             db[csv_table].transform(types=tracker.types)

sqlite_utils/utils.py

@@ -513,3 +513,21 @@ def hash_record(record: Dict, keys: Optional[Iterable[str]] = None):
             "utf8"
         )
     ).hexdigest()
+
+
+def _flatten(d):
+    for key, value in d.items():
+        if isinstance(value, dict):
+            for key2, value2 in _flatten(value):
+                yield key + "_" + key2, value2
+        else:
+            yield key, value
+
+
+def flatten(row: dict) -> dict:
+    """
+    Turn a nested dict e.g. ``{"a": {"b": 1}}`` into a flat dict: ``{"a_b": 1}``
+
+    :param row: A Python dictionary, optionally with nested dictionaries
+    """
+    return dict(_flatten(row))

tests/test_cli.py

@@ -2176,18 +2176,6 @@ def test_upsert_detect_types(tmpdir, option):
     ]
 
 
-@pytest.mark.parametrize(
-    "input,expected",
-    (
-        ({"foo": {"bar": 1}}, {"foo_bar": 1}),
-        ({"foo": {"bar": [1, 2, {"baz": 3}]}}, {"foo_bar": [1, 2, {"baz": 3}]}),
-        ({"foo": {"bar": 1, "baz": {"three": 3}}}, {"foo_bar": 1, "foo_baz_three": 3}),
-    ),
-)
-def test_flatten_helper(input, expected):
-    assert dict(cli._flatten(input)) == expected
-
-
 def test_integer_overflow_error(tmpdir):
     db_path = str(tmpdir / "test.db")
     result = CliRunner().invoke(

tests/test_utils.py

@@ -71,3 +71,15 @@ def test_maximize_csv_field_size_limit():
     assert len(rows_list2) == 1
     assert rows_list2[0]["id"] == "1"
     assert rows_list2[0]["text"] == long_value
+
+
+@pytest.mark.parametrize(
+    "input,expected",
+    (
+        ({"foo": {"bar": 1}}, {"foo_bar": 1}),
+        ({"foo": {"bar": [1, 2, {"baz": 3}]}}, {"foo_bar": [1, 2, {"baz": 3}]}),
+        ({"foo": {"bar": 1, "baz": {"three": 3}}}, {"foo_bar": 1, "foo_baz_three": 3}),
+    ),
+)
+def test_flatten(input, expected):
+    assert utils.flatten(input) == expected