Coverage for C:\src\imod-python\imod\util\nested_dict.py: 94%

31 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1import collections 

2import functools 

3from typing import Any, Dict, List 

4 

5 

6def initialize_nested_dict(depth: int) -> collections.defaultdict: 

7 """ 

8 Initialize a nested dict with a fixed depth 

9 

10 Parameters 

11 ---------- 

12 depth : int 

13 depth of returned nested dict 

14 

15 Returns 

16 ------- 

17 nested defaultdicts of n depth 

18 

19 """ 

20 # In explicit form, say we have ndims=5 

21 # Then, writing it out, we get: 

22 # a = partial(defaultdict, {}) 

23 # b = partial(defaultdict, a) 

24 # c = partial(defaultdict, b) 

25 # d = defaultdict(c) 

26 # This can obviously be done iteratively. 

27 if depth == 0: 

28 return {} 

29 elif depth == 1: 

30 return collections.defaultdict(dict) 

31 else: 

32 d = functools.partial(collections.defaultdict, dict) 

33 for _ in range(depth - 2): 

34 d = functools.partial(collections.defaultdict, d) 

35 return collections.defaultdict(d) 

36 

37 

38def set_nested(d: collections.defaultdict, keys: List[str], value: Any) -> None: 

39 """ 

40 Set in the deepest dict of a set of nested dictionaries, as created by the 

41 initialize_nested_dict function above. 

42 

43 Mutates d. 

44 

45 Parameters 

46 ---------- 

47 d : (Nested dict of) dict 

48 keys : list of keys 

49 Each key is a level of nesting 

50 value : dask array, typically 

51 

52 Returns 

53 ------- 

54 None 

55 """ 

56 if len(keys) == 1: 

57 d[keys[0]] = value 

58 else: 

59 set_nested(d[keys[0]], keys[1:], value) 

60 

61 

62def append_nested_dict(dict1: Dict, dict2: Dict) -> None: 

63 """ 

64 Recursively walk through two dicts to append dict2 to dict1. 

65 

66 Mutates dict1 

67 

68 Modified from: 

69 https://stackoverflow.com/a/58742155 

70 

71 Parameters 

72 ---------- 

73 dict1 : nested dict 

74 Nested dict to be appended to 

75 dict2 : nested dict 

76 Nested dict to append 

77 

78 """ 

79 for key, val in dict1.items(): 

80 if isinstance(val, dict): 

81 if key in dict2 and isinstance(dict2[key], dict): 

82 append_nested_dict(dict1[key], dict2[key]) 

83 else: 

84 if key in dict2: 

85 dict1[key] = dict2[key] 

86 

87 for key, val in dict2.items(): 

88 if key not in dict1: 

89 dict1[key] = val 

90 

91 

92def sorted_nested_dict(d: Dict) -> Dict: 

93 """ 

94 Sorts a variably nested dict (of dicts) by keys. 

95 

96 Each dictionary will be sorted by its keys. 

97 

98 Parameters 

99 ---------- 

100 d : (Nested dict of) dict 

101 

102 Returns 

103 ------- 

104 sorted_lists : list (of lists) 

105 Values sorted by keys, matches the nesting of d. 

106 """ 

107 firstkey = next(iter(d.keys())) 

108 if not isinstance(d[firstkey], dict): # Base case 

109 return [v for (_, v) in sorted(d.items(), key=lambda t: t[0])] 

110 else: # Recursive case 

111 return [ 

112 sorted_nested_dict(v) for (_, v) in sorted(d.items(), key=lambda t: t[0]) 

113 ]