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

1import pandas as pd 

2import plotly.graph_objects as go 

3 

4 

5def _hover_breakdown(transactions, node_name, hover_field): 

6 """ 

7 Build hover text for a node: average spend/day plus a breakdown by hover_field, for all 

8 transactions flowing into that node. Returns "" if there's nothing to break down (eg. a 

9 node that's never a Target). 

10 """ 

11 df = transactions.processed_data 

12 node_rows = df[df['Target'] == node_name] 

13 breakdown = node_rows.groupby([hover_field]).agg({'Amount': 'sum'}) 

14 if len(breakdown) == 0: 

15 return "" 

16 days = (df['Date'].max() - df['Date'].min()).days 

17 avg_per_day = node_rows['Amount'].sum() / days if days else 0 

18 text = f"Avg/day: ${avg_per_day:.2f}<br>-------------------<br>Categories:<br>" 

19 for item, amount in breakdown['Amount'].items(): 

20 text += f"{item}: ${amount:.2f}<br>" 

21 return text 

22 

23 

24def build_sankey_figure(transactions, labels, app_settings) -> go.Figure: 

25 """ 

26 Build a plotly Sankey figure from processed transaction data. `transactions` must have 

27 already been through Transactions.process(). 

28 """ 

29 grouped = transactions.grouped_data.copy() # Work on a copy - this function must be safe to call more than once. 

30 unique_nodes = list(pd.unique(grouped[['Source', 'Target']].values.ravel('K'))) 

31 node_index = {name: idx for idx, name in enumerate(unique_nodes)} 

32 

33 node_colors = [labels.get_attribute(name, "node_color") for name in unique_nodes] 

34 

35 # Link color: prefer the target node's color, falling back to the source node's (or the 

36 # sheet's default color if neither is set). 

37 link_colors = [] 

38 for source, target in zip(grouped['Source'], grouped['Target']): 

39 color = labels.get_attribute(target, "link_color", use_default=False) 

40 if not color: 

41 color = labels.get_attribute(source, "link_color") 

42 link_colors.append(color) 

43 

44 node_settings = { 

45 'pad': 15, 

46 'thickness': 20, 

47 'line': dict(color='black', width=0.5), 

48 'label': unique_nodes, 

49 'color': node_colors, 

50 } 

51 

52 if app_settings.hover: 

53 node_settings['customdata'] = [ 

54 _hover_breakdown(transactions, name, app_settings.hover) for name in unique_nodes 

55 ] 

56 node_settings['hovertemplate'] = 'Total: %{value}<br>%{customdata}<extra></extra>' 

57 

58 fig = go.Figure(data=[go.Sankey( 

59 valueformat="$.2f", 

60 node=node_settings, 

61 link=dict( 

62 source=grouped['Source'].map(node_index), 

63 target=grouped['Target'].map(node_index), 

64 value=grouped['Amount'], 

65 color=link_colors, 

66 ) 

67 )]) 

68 

69 title = go.layout.Title({'font': {'family': 'Courier New', 'size': 12}, 'text': transactions.title}) 

70 fig.update_layout(title=title) 

71 return fig 

72 

73 

74def _sum_row_amount(row): 

75 """ 

76 Sum Amount + Sales Tax + Tips for a row, treating unparseable/missing values as zero. 

77 """ 

78 total = 0 

79 for col in ("Amount", "Sales Tax", "Tips"): 

80 try: 

81 total += float(row[col]) 

82 except (ValueError, TypeError): 

83 pass 

84 return total 

85 

86 

87def build_line_figure(transactions, app_settings) -> go.Figure: 

88 """ 

89 Build a plotly line chart of spend over time, grouped by Classification. 

90 

91 NOTE: less mature than build_sankey_figure() - can get noisy for large datasets, doesn't 

92 fill in zero-value gaps for categories with sparse activity, and doesn't support a log 

93 scale. 'Income', 'Uncategorized', and any classification prefixed with 'x' (a convention 

94 used in sample_data/labels.csv, eg. 'xEntertainment') are hidden by default. 

95 """ 

96 df = transactions.processed_data.assign(**{"Total Amount": None}) 

97 df["Total Amount"] = df.apply(_sum_row_amount, axis=1) 

98 

99 classifications = sorted(df["Classification"].unique()) 

100 date_idx = pd.date_range(start=df["Date"].min(), end=df["Date"].max()) 

101 

102 fig = go.Figure() 

103 for classification in classifications: 

104 if classification in ("Income", "Uncategorized"): 

105 continue 

106 if classification.startswith('x'): 

107 continue 

108 series = df[df["Classification"] == classification].groupby('Date')["Total Amount"].sum() 

109 series.index = pd.DatetimeIndex(series.index) 

110 series = series.reindex(date_idx, fill_value=float("nan")) # connectgaps needs NaN, not 0 

111 if app_settings.chart_resolution == 'day': 

112 series = series.resample('D', label='left').sum() 

113 elif app_settings.chart_resolution == 'week': 

114 series = series.resample('W', label='left').sum() 

115 elif app_settings.chart_resolution == 'month': 

116 series = series.resample('ME', label='left').sum() 

117 series = series.replace(0, float("nan")) 

118 fig.add_trace(go.Scatter( 

119 y=series.to_list(), x=series.index.to_list(), mode='lines', name=classification, connectgaps=True 

120 )) 

121 return fig