diff --git a/src/coola/nested/__init__.py b/src/coola/nested/__init__.py new file mode 100644 index 00000000..290e178d --- /dev/null +++ b/src/coola/nested/__init__.py @@ -0,0 +1,10 @@ +r"""Contain simple operations on nested data structure.""" + +from __future__ import annotations + +__all__ = [ + "convert_to_dict_of_lists", + "convert_to_list_of_dicts", +] + +from coola.nested.conversion import convert_to_dict_of_lists, convert_to_list_of_dicts diff --git a/src/coola/nested/conversion.py b/src/coola/nested/conversion.py new file mode 100644 index 00000000..5ef7058d --- /dev/null +++ b/src/coola/nested/conversion.py @@ -0,0 +1,66 @@ +r"""Contain some utility functions to convert nested data structure.""" + +from __future__ import annotations + +__all__ = [ + "convert_to_dict_of_lists", + "convert_to_list_of_dicts", +] + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Hashable, Mapping, Sequence + + +def convert_to_dict_of_lists(seq_of_mappings: Sequence[Mapping]) -> dict[Hashable, list]: + r"""Convert a sequence of mappings to a dictionary of lists. + + All the dictionaries should have the same keys. The first + mapping in the sequence is used to find the keys. + + Args: + seq_of_mappings: The sequence of mappings to convert. + + Returns: + A dictionary of lists. + + Example usage: + + ```pycon + + >>> from coola.nested import convert_to_dict_of_lists + >>> convert_to_dict_of_lists( + ... [{"key1": 1, "key2": 10}, {"key1": 2, "key2": 20}, {"key1": 3, "key2": 30}] + ... ) + {'key1': [1, 2, 3], 'key2': [10, 20, 30]} + + ``` + """ + if seq_of_mappings: + return {key: [dic[key] for dic in seq_of_mappings] for key in seq_of_mappings[0]} + return {} + + +def convert_to_list_of_dicts(mapping_of_seqs: Mapping[Hashable, Sequence]) -> list[dict]: + r"""Convert a mapping of sequences to a list of dictionaries. + + All the sequences should have the same length. + + Args: + mapping_of_seqs: The mapping of sequences to convert. + + Returns: + A dictionary of lists. + + Example usage: + + ```pycon + + >>> from coola.nested import convert_to_list_of_dicts + >>> convert_to_list_of_dicts({"key1": [1, 2, 3], "key2": [10, 20, 30]}) + [{'key1': 1, 'key2': 10}, {'key1': 2, 'key2': 20}, {'key1': 3, 'key2': 30}] + + ``` + """ + return [dict(zip(mapping_of_seqs, seqs)) for seqs in zip(*mapping_of_seqs.values())] diff --git a/tests/unit/nested/__init__.py b/tests/unit/nested/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/nested/test_conversion.py b/tests/unit/nested/test_conversion.py new file mode 100644 index 00000000..ccdf7a9b --- /dev/null +++ b/tests/unit/nested/test_conversion.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from coola.nested import convert_to_dict_of_lists, convert_to_list_of_dicts + +############################################## +# Tests for convert_to_dict_of_lists # +############################################## + + +def test_convert_to_dict_of_lists_empty_list() -> None: + assert convert_to_dict_of_lists([]) == {} + + +def test_convert_to_dict_of_lists_empty_dict() -> None: + assert convert_to_dict_of_lists([{}]) == {} + + +def test_convert_to_dict_of_lists() -> None: + assert convert_to_dict_of_lists( + [{"key1": 1, "key2": 10}, {"key1": 2, "key2": 20}, {"key1": 3, "key2": 30}] + ) == { + "key1": [1, 2, 3], + "key2": [10, 20, 30], + } + + +############################################## +# Tests for convert_to_list_of_dicts # +############################################## + + +def test_convert_to_list_of_dicts_empty_dict() -> None: + assert convert_to_list_of_dicts({}) == [] + + +def test_convert_to_list_of_dicts_empty_list() -> None: + assert convert_to_list_of_dicts({"key1": [], "key2": []}) == [] + + +def test_convert_to_list_of_dicts() -> None: + assert convert_to_list_of_dicts({"key1": [1, 2, 3], "key2": [10, 20, 30]}) == [ + {"key1": 1, "key2": 10}, + {"key1": 2, "key2": 20}, + {"key1": 3, "key2": 30}, + ]