Coverage for src/sankey_cashflow/labels.py: 95%

1from typing import Optional, Union 

2 

3import networkx as nx 

4from numpy.typing import ArrayLike 

5 

6from .utils import is_null, logger 

7 

8 

9class RowLabels: 

10 """ 

11 Contains row label data and methods, used to map data categories to source, target, and color information. 

12 Accepts an array of dicts on init which can come from a Google sheet or csv. 

13 Data looks like: [{'Category Name': 'House', 'Type': 'computed', 'Source': 'Income', 'Target': 'House', 

14 'Link color': 'rgba(153, 187, 255, 0.8)', 'Node color': 'rgba(102, 153, 255, 1)', 'Comments': '', '': ''}] 

15 Also creates a DAG with source target edges based on labels definitions 

16 """ 

17 def __init__(self, labeldata: list[dict]): 

18 self._labeldata = labeldata 

19 self._digraph = nx.DiGraph() 

20 self._digraph.add_node("Income", ntype="income") 

21 self.process_report = f"RowLabel report\n{'=' * 60}\n\n\n" 

22 required_columns = ['Category Name', 'Type', 'Source', 'Target', 'Classification', 'Link color', 'Node color'] 

23 if False in [k in self._labeldata[0].keys() for k in required_columns]: 

24 raise Exception(f"Sources-Targets sheet does not have all required columns! Needed: {required_columns}. \ 

25 Found: {list(self._labeldata[0].keys())}") 

26 self._available_attributes = ['source', 'target', 'classification', 'link_color', 'node_color', 'type'] 

27 # TODO: validation 

28 self._lookup = {} 

29 # Map out data to lookup dict 

30 for i in self._labeldata: 

31 if len(i['Category Name']) == 0: # Skip empties 

32 self.process_report += f"SKIPPING: {i}\n" 

33 continue 

34 if i['Category Name'] in self._lookup: 

35 raise Exception(f"Duplicate label! {i['Category Name']}") 

36 # Add to internal lookup dict 

37 self.process_report += f"ADDING LOOKUP: {i}\n" 

38 self._lookup[i['Category Name']] = { 

39 'source': i.get('Source'), 

40 'target': i.get('Target'), 

41 'classification': i.get('Classification'), 

42 'link_color': i.get('Link color'), 

43 'node_color': i.get('Node color'), 

44 'type': i.get('Type') 

45 } 

46 if i.get("Source") == "DEDUCTIONS": 

47 # These have to be dynamically generated - skip adding to DAG for now (or maybe entirely) 

48 # TODO (maybe): switch to using type and/or DAG attributes 

49 self.process_report += f"SKIPPING DEDUCTION: {i}\n" 

50 continue 

51 if i["Type"] in ['tag', 's-tag']: 

52 # NOT adding tag labels to graph - if needed they will be added on the fly 

53 # TODO: investigate using DAG attributes instead. Test with distant/complex tag targets. 

54 if i["Type"] == 's-tag': 

55 self.process_report += f"Adding NODE for s-tag: {i} to DAG\n" 

56 self._digraph.add_node(i['Category Name'], type='s-tag') 

57 else: 

58 self.process_report += f"NOT Adding tag: {i} to DAG\n" 

59 continue 

60 if i.get('Source') and i.get('Target'): # Add all source/target pairs as edges to DAG 

61 # Add to DAG 

62 if i.get('Type'): 

63 self.process_report += \ 

64 f"ADDING EDGE: {i.get('Source')} -> {i.get('Target')} (ntype={i.get('Type')})\n" 

65 self._digraph.add_edge(i.get('Source'), i.get('Target'), ntype=i.get('Type')) 

66 else: 

67 self.process_report += f"ADDING EDGE: {i.get('Source')} -> {i.get('Target')}\n" 

68 self._digraph.add_edge(i.get('Source'), i.get('Target')) 

69 else: 

70 self.process_report += f"ERROR (SKIPPING): {i}\n" 

71 logger.warning(f"Category: {i['Category Name']} yielded an empty source and/or target! \ 

72 {i.get('Source')}:{i.get('Target')}") 

73 

74 @property 

75 def data(self) -> list[dict]: 

76 return self._labeldata 

77 

78 @property 

79 def graph(self) -> nx.DiGraph: 

80 return self._digraph 

81 

82 def get_longest_path(self) -> ArrayLike: 

83 return nx.dag_longest_path(self._digraph) 

84 

85 def get_path(self, source: str, target: str) -> Union[list, None]: 

86 path = [i for i in nx.all_simple_paths(self._digraph, source, target)] 

87 # Note path obj may contain 0 or multiple paths 

88 if not path or len(path) == 0: 

89 logger.warning(f"No path found for {source}:{target}") 

90 elif len(path) > 1: 

91 logger.warning(f"Multiple paths found for {source}:{target}. ({path})") 

92 else: 

93 return path[0] 

94 

95 def get_label(self, labelname: str, labeltype: Union[str, None] = None) -> Union[tuple[str, str], None]: 

96 """ 

97 Return a label name & tag pair or None 

98 """ 

99 item = self._lookup.get(labelname) # TODO: test duplicate category names 

100 if labeltype == "any": 

101 return item 

102 if labeltype and labeltype == "tag": 

103 if item and item.get('type') == 'tag': 

104 return item 

105 return None 

106 elif labeltype and labeltype == "s-tag": 

107 if item and item.get('type') == 's-tag': 

108 return item 

109 return None 

110 if item and item.get('type') not in ['tag', 's-tag']: 

111 return item 

112 return None 

113 

114 def get_attribute(self, 

115 labelname: str, 

116 labelattribute: str, 

117 use_default: Optional[bool] = True, 

118 labeltype: Optional[Union[str, None]] = None, 

119 original_category: Optional[Union[str, None]] = None) -> Union[str, None]: 

120 """ 

121 Get named attribute by label. Cases: 

122 Unknown attribute: raise exception 

123 label & labeltype are exist: 

124 attribute exists: return attribute 

125 attribute doesn't exist: return default attribute 

126 label & labeltype doesn't exists: 

127 return default attributes 

128 """ 

129 if labelattribute not in self._available_attributes: 

130 raise Exception(f"Unknown attribute: {labelattribute}") 

131 if labelattribute == "type": 

132 item = self.get_label(labelname, "any") 

133 else: 

134 item = self.get_label(labelname, labeltype) 

135 default_item = self.get_label('default') 

136 if item: 

137 if labelattribute == "type": 

138 return item.get("type", "") 

139 if not is_null(item[labelattribute]) and len(item[labelattribute]) != 0: 

140 return item[labelattribute] 

141 elif labelattribute == "type": 

142 return None 

143 if use_default: 

144 # Did not find the attribute - use defaults 

145 if labelattribute == "source": 

146 return "Uncategorized" 

147 if labelattribute == "target": 

148 return labelname 

149 if labelattribute == "classification": 

150 return 'Uncategorized' 

151 if labelattribute in ["node_color", "link_color"]: 

152 return default_item[labelattribute] 

153 return None 

154 

155 def print_graph(self, filename: str) -> None: 

156 from matplotlib import pyplot as plt 

157 plt.tight_layout() 

158 nx.draw_networkx(self._digraph, arrows=True) 

159 # plt.figure(figsize=(20,20)) 

160 plt.savefig(filename, dpi=200, format="PNG") 

161 plt.clf()