Coverage for src/sankey_cashflow/cli.py: 72%

1import argparse 

2import datetime 

3 

4import pandas as pd 

5 

6from .diagram import build_line_figure, build_sankey_figure 

7from .io import fetch_data, read_csv_as_df 

8from .labels import RowLabels 

9from .settings import AppSettings 

10from .transactions import Transactions 

11from .utils import logger, save_report, validate_date_string 

12 

13 

14def build_arg_parser() -> argparse.ArgumentParser: 

15 parser = argparse.ArgumentParser(prog="sankeyd", description="Generate a cashflow Sankey or line diagram.") 

16 parser.add_argument("-s", "--source", help="Data source, either a Google sheet or a .csv") 

17 parser.add_argument("-t", "--sheet", help="Sheet name, defaults to 'current_exp'") 

18 parser.add_argument("--srcmap", help="Location for sources-targets mapping csv (only needed for csv data " 

19 "sources). Defaults to 'Sources-Targets' for GSheets") 

20 parser.add_argument("-r", "--range", help="Enable time range filtering", action="store_true") 

21 parser.add_argument("--separate_tax", help="Show taxes as own flow", action="store_true") 

22 parser.add_argument("-v", "--verbose", help="Enable verbose output", action="store_true") 

23 # Note: the specified service account needs to have access to the data source document. 

24 parser.add_argument("--creds", help="Google service account credentials to use (in json format).") 

25 parser.add_argument("--distributions", help="Distribute transactions where noted", action="store_true") 

26 parser.add_argument("--tags", help="Comma delimited list of tags to search for, eg 'foo, bar, baz bat'.") 

27 parser.add_argument("--exclude", help="Comma delimited list of tags to exclude, eg 'foo, bar, baz bat'.") 

28 parser.add_argument("--tag_override", action="store_true", 

29 help="Use tag labels to override source/targets instead of appending.") 

30 parser.add_argument("--stores", help="Comma delimited list of stores/payees to override source-targets on, " 

31 "eg 'Safeway, Bartells, Fred Meyer'.") 

32 parser.add_argument("--all_time", action="store_true", 

33 help="Include transactions that happen in the future and before 10/1/2022") 

34 parser.add_argument("--feed_in", help="If using source-tags, feed surplus into income.", action="store_true") 

35 parser.add_argument("--audit", help="Audit source data transactions for missing items", action="store_true") 

36 parser.add_argument("--recurring", help="Split recurring expenses out", action="store_true") 

37 parser.add_argument("--hover", help="Grouping field for hover text. Defaults to 'Category'") 

38 parser.add_argument("--dtype", help="Diagram type to generate (sankey, line). Defaults to 'sankey'") 

39 return parser 

40 

41 

42def _prompt_for_date_range(app_settings) -> tuple: 

43 sdate = input("Enter start date, in form of YYYY-MM-DD or MM/DD/YYYY. (<Enter> to leave unbounded): ") 

44 while not validate_date_string(sdate, True): 

45 sdate = input(f"Invalid date string ({sdate})! Please enter start date, in form of YYYY-MM-DD or " 

46 "MM/DD/YYYY. (<Enter> to leave unbounded): ") 

47 edate = input("Enter end date, in form of YYYY-MM-DD. (<Enter> to leave unbounded): ") 

48 if not edate: 

49 edate = datetime.datetime.now().strftime("%Y-%m-%d") 

50 while not validate_date_string(edate, True): 

51 edate = input(f"Invalid date string ({edate})! Please enter end date, in form of YYYY-MM-DD or " 

52 "MM/DD/YYYY. (<Enter> to leave unbounded): ") 

53 try: 

54 app_settings.date_filter_start = sdate 

55 app_settings.date_filter_end = edate 

56 if app_settings.date_filter_start is None and app_settings.date_filter_end is None: 

57 print("Need to have at least one of start date or end date to use date filtering!") 

58 raise SystemExit(1) 

59 except Exception as e: 

60 print(f"Could not parse supplied date(s). Error was: {e}") 

61 raise SystemExit(1) 

62 return sdate, edate 

63 

64 

65def main(argv=None): 

66 args = build_arg_parser().parse_args(argv) 

67 app_settings = AppSettings(args) 

68 

69 date_range = None 

70 if app_settings.filter_dates: 

71 # Doing this first, as it involves user input - throw any errors related to that before fetching data. 

72 date_range = _prompt_for_date_range(app_settings) 

73 

74 if app_settings.diagram_type == 'line': 

75 chart_resolution = input("Enter chart resolution (day, week, month): ") 

76 while chart_resolution not in ['day', 'week', 'month']: 

77 chart_resolution = input("Invalid chart resolution! Enter chart resolution (day, week, month): ") 

78 app_settings.chart_resolution = chart_resolution 

79 

80 if app_settings.verbose: 

81 logger.info(f"Fetching data from {app_settings.data_source}: {app_settings.data_sheet}...") 

82 src_target, df = fetch_data(app_settings) 

83 

84 sources_targets = RowLabels(src_target) 

85 transactions_data = Transactions(df, sources_targets, app_settings) 

86 

87 save_report(sources_targets.process_report, 'labels_report') 

88 

89 if app_settings.audit_mode: 

90 print("We are in audit mode, so we will not process the data.") 

91 audit_file = input("Enter the path to the bank transactions file: ") 

92 audit_data = read_csv_as_df(audit_file) 

93 audit_data["Date"] = pd.to_datetime(audit_data["Date"]) # Normalize date format 

94 audit_data["Amount"] = audit_data["Amount"].apply(lambda x: x * -1) # Convert from negative to positive 

95 audit_report = transactions_data.audit(audit_data, date_range) 

96 print(audit_report) 

97 raise SystemExit 

98 

99 if app_settings.diagram_type == 'sankey': 

100 transactions_data.process(date_range) 

101 elif app_settings.diagram_type == 'line': 

102 transactions_data.process_line(date_range) 

103 else: 

104 print(f"Invalid diagram type: {app_settings.diagram_type}") 

105 raise SystemExit(1) 

106 

107 save_report(transactions_data.process_report, 'transactions_report') 

108 

109 if app_settings.verbose: 

110 logger.info(f"Generating diagram of type {app_settings.diagram_type}") 

111 

112 if app_settings.diagram_type == 'sankey': 

113 fig = build_sankey_figure(transactions_data, sources_targets, app_settings) 

114 else: 

115 fig = build_line_figure(transactions_data, app_settings) 

116 fig.show() 

117 

118 

119if __name__ == '__main__': 

120 main()