Coverage for src/io_collection/keys/group_keys.py: 100%

1from itertools import groupby 

2 

3 

4def group_keys(keys: list[str]) -> dict[str, list[str]]: 

5 """ 

6 Group keys based on individual key parts. 

7 

8 A key is composed of parts separated by a single underscores (`_`). Keys are 

9 grouped such that each key group consists of the subset of keys with the 

10 matching part at the given position. 

11 

12 Returns 

13 ------- 

14 : 

15 The map of key groups to keys. 

16 """ 

17 

18 parts = [key.split("_") for key in keys] 

19 num_parts = len(parts[0]) 

20 

21 if not all(len(part) == num_parts for part in parts): 

22 message = "All keys must have the same number of parts" 

23 raise ValueError(message) 

24 

25 groups = {} 

26 

27 for group_index in range(num_parts): 

28 sorted_keys = sorted(parts, key=lambda k: k[group_index]) 

29 for group, part in groupby(sorted_keys, key=lambda k: k[group_index]): 

30 groups[group] = ["_".join(p) for p in part] 

31 

32 return groups