|
| 1 | +def flatten(data: dict | list[dict]) -> dict | list[dict]: |
| 2 | + """Flatten a nested dictionary into a single-level dictionary. |
| 3 | +
|
| 4 | + Parameters |
| 5 | + ---------- |
| 6 | + data |
| 7 | + A dictionary or a list of dictionaries to flatten. |
| 8 | +
|
| 9 | + Returns |
| 10 | + ------- |
| 11 | + If `data` is a dictionary, returns a flattened dictionary. |
| 12 | + If `data` is a list of dictionaries, returns a list of flattened dictionaries. |
| 13 | + """ |
| 14 | + if isinstance(data, dict): |
| 15 | + return _flatten(data) |
| 16 | + if isinstance(data, list): |
| 17 | + if not data: |
| 18 | + return [] |
| 19 | + if not all(isinstance(item, dict) for item in data): |
| 20 | + raise ValueError("All items in the list must be dictionaries.") |
| 21 | + return [_flatten(item) for item in data] |
| 22 | + raise ValueError("Input must be a dictionary or a list of dictionaries.") |
| 23 | + |
| 24 | + |
| 25 | +def _flatten(data, output: dict | None = None, path: str | None = None): |
| 26 | + output = output if output is not None else {} |
| 27 | + if isinstance(data, str | int | float | bool | None): |
| 28 | + output[path] = data |
| 29 | + return output |
| 30 | + if isinstance(data, dict): |
| 31 | + for key, value in data.items(): |
| 32 | + new_path = _flatten_key(path, key) |
| 33 | + _flatten(data=value, output=output, path=new_path) |
| 34 | + return output |
| 35 | + if isinstance(data, list): |
| 36 | + if len(data) == 0: |
| 37 | + output[path] = None |
| 38 | + return output |
| 39 | + if all(isinstance(val, str | int | float | bool | None) for val in data): |
| 40 | + output[path] = data[0] if len(data) == 1 else data |
| 41 | + return output |
| 42 | + if len(data) != 1: |
| 43 | + raise ValueError(f"Data at {path} is a list with more than one non-primitive element: {data}") |
| 44 | + return _flatten(data=data[0], output=output, path=path) |
| 45 | + raise ValueError(f"Unknown data type {type(data)} at {path}") |
| 46 | + |
| 47 | + |
| 48 | +def _flatten_key(path: str | None, key: str): |
| 49 | + if path: |
| 50 | + return f"{path}.{key}" |
| 51 | + return key |
0 commit comments