diff --git a/docs/reference.rst b/docs/reference.rst index 66b9d38..5b5fd25 100644 --- a/docs/reference.rst +++ b/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 diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index bf7fae4..5fcc95a 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -24,6 +24,7 @@ from .utils import ( 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) diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 8754554..8916c46 100644 --- a/sqlite_utils/utils.py +++ b/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)) diff --git a/tests/test_cli.py b/tests/test_cli.py index 24f36ed..b8df8d6 100644 --- a/tests/test_cli.py +++ b/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( diff --git a/tests/test_utils.py b/tests/test_utils.py index d397176..f728bcd 100644 --- a/tests/test_utils.py +++ b/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