Coverage for src/sankey_cashflow/io.py: 44%
109 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-05 04:36 +0000
1from csv import DictReader
2from os import path
3from pathlib import Path
4from typing import Union
6import pandas as pd
7import pygsheets
9from .data_row import DataRow
10from .utils import logger, normalize_amounts
13def fetch_data(app_settings_obj): # source_spreadsheet, source_worksheet, csv_src_target, service_account_credentials):
14 # app_settings_obj.data_source,app_settings_obj.data_sheet,app_settings_obj.labels_source,app_settings_obj.g_creds
15 # WIP: Handle source data in mutiple sheets, eg. 1 sheet per year...
17 """
18 self.data_source = args.source # A csv file or a google Sheets document containing transactions data.
19 (In the case of the latter, a sheet name must be provided as well)
20 self.data_sheet = args.sheet or "Transactions_*" # The name (or prefix plus wildcard) of the sheet containing
21 the transactions data (if data_source is a google Sheets document
22 self._labels_source = args.srcmap or "Sources-Targets" #A csv file or sheet name containing sources and targets
23 """
24 def data_source_router(
25 filename: str,
26 file_kind: str,
27 gcreds,
28 sheetname: Union[str, None] = None,
29 gcreds_obj=None,
30 gsheets_obj=None):
31 """
32 Calls itself recursively, looking for data sources either locally or in Google Sheets and returns the first hit.
33 Handle wildcards expressions for multiple sheets.
34 Return a dict containing a list of filenames.
35 :param filename: A filename or wildcard expression, or None (Note: only csv files will use the wildcard
36 expression for this value)
37 :param file_kind: "sources-targets" or "transactions"
38 :param gcreds: Google credentials object
39 :param sheetname: A sheet name or wildcard expression, or None (only applicable to Google Sheets)
40 :param gc: Google client object (optional)
42 Returns: {
43 "filename" -> list of strings,
44 "sheetname" -> list of strings,
45 "filetype" -> ["csv", "gsheet"],
46 "file_kind" -> ["sources-targets","transactions"],
47 "gcreds_obj" -> Google credentials object",
48 "gsheets_obj" -> Google sheets object
49 }
50 """
51 file_kind = file_kind.lower()
52 if file_kind not in ["sources-targets", "transactions"]:
53 raise Exception(f"Invalid file kind: {file_kind}")
54 if filename is None:
55 filename = input(f"Enter location for {file_kind} data: ")
56 return data_source_router(filename, file_kind, gcreds, sheetname, gcreds_obj, gsheets_obj)
57 if filename.endswith(".csv"):
58 if path.isfile(filename): # Switch to pathlib?
59 return {
60 "filename": [filename],
61 "sheetname": [None],
62 "filetype": "csv",
63 "file_kind": file_kind,
64 "gcreds_obj": gcreds_obj,
65 "gsheets_obj": gsheets_obj}
66 logger.warning(f"File not found: {filename}")
67 return data_source_router(None, file_kind, gcreds, sheetname, gcreds_obj, gsheets_obj)
68 if filename.endswith("*"):
69 # Wildcards on filenames only supported for csvs, Gsheets would use sheetnames for wildcard.
70 fileparts = filename.split("/")
71 filepattern = f"{fileparts[-1]}.csv"
72 if len(fileparts[0]) == 0:
73 # This must be an absolute path
74 this_dir = "/" + "/".join(fileparts[1:-1])
75 else:
76 # Possibly a relative path
77 if len(fileparts) > 1:
78 this_dir = "./" + "/".join(fileparts[:-1])
79 else:
80 this_dir = "."
81 dir_obj = Path(this_dir)
82 if dir_obj.is_dir():
83 files = list(dir_obj.glob(filepattern))
84 if len(files) > 0:
85 return {"filename": [str(f) for f in files],
86 "sheetname": [None],
87 "filetype": "csv",
88 "file_kind": file_kind,
89 "gcreds_obj": gcreds_obj,
90 "gsheets_obj": gsheets_obj}
91 logger.warning("No files found at: {filename}")
92 return data_source_router(None, file_kind, gcreds, sheetname, gcreds_obj, gsheets_obj)
93 # If we've gotten this far, we must be looking for a Google Sheets file (late binding of the
94 # gc object to reduce calls to authorize)
95 if not gcreds_obj:
96 gcreds_obj = pygsheets.authorize(service_file=gcreds) # trying to avoid calling this multiple times
97 if filename not in gcreds_obj.spreadsheet_titles():
98 logger.warning(f"Spreadsheet not found: {filename}")
99 return data_source_router(None, file_kind, gcreds, sheetname, gcreds_obj, gsheets_obj)
100 gsheets_obj = gcreds_obj.open(filename) # TODO: error handling.
101 if not sheetname:
102 sheetname = input(f"Enter worksheet title for spreadsheet {filename}: ")
103 return data_source_router(filename, file_kind, gcreds, sheetname, gcreds_obj, gsheets_obj)
104 gsheet_titles = [i.title for i in gsheets_obj.worksheets()]
105 if sheetname.endswith("*"):
106 # Wildcard handling for sheetnames
107 results = []
108 for sheet in gsheet_titles:
109 if sheet.startswith(sheetname[:-1]):
110 results.append(sheet)
111 if len(results) > 0:
112 return {
113 "filename": [filename],
114 "sheetname": results,
115 "filetype": "gsheet",
116 "file_kind": file_kind,
117 "gcreds_obj": gcreds_obj,
118 "gsheets_obj": gsheets_obj}
119 logger.warning(f"No spreadsheets found matching pattern: \"{sheetname}\" in Google Sheets: \"{filename}\"")
120 return data_source_router(filename, file_kind, gcreds, None, gcreds_obj, gsheets_obj)
121 if sheetname in gsheet_titles:
122 return {"filename": [filename],
123 "sheetname": [sheetname],
124 "filetype": "gsheet",
125 "file_kind": file_kind,
126 "gcreds_obj": gcreds_obj,
127 "gsheets_obj": gsheets_obj}
128 logger.warning(f"Spreadsheet \"{sheetname}\" not found in Google Sheets: \"{filename}\"")
129 return data_source_router(filename, file_kind, gcreds, None, gcreds_obj, gsheets_obj)
131 try:
132 # Get sources-targets file location data
133 # See sample_data/expenses.csv and labels.csv for examples of data format.
134 # TODO: validation & error handling
135 # TODO: there should only ever be one sources-targets file, so data_source_router should handle that
136 src_target = None
137 if app_settings_obj.labels_source.endswith(".csv"):
138 src_target_file = data_source_router(app_settings_obj.labels_source,
139 "sources-targets",
140 app_settings_obj.g_creds, None, None, None)
141 else:
142 src_target_file = data_source_router(app_settings_obj.data_source,
143 "sources-targets",
144 app_settings_obj.g_creds,
145 app_settings_obj.labels_source, None, None)
146 if src_target_file["filetype"] == "csv":
147 # Sources-targets data is returned differently than transactions data, so we need to handle it differently.
148 with open(src_target_file["filename"][0], 'r') as f:
149 dr = DictReader(f)
150 src_target = list(dr)
151 else:
152 # Fetch source-target info and colors.
153 src_target = src_target_file["gsheets_obj"].\
154 worksheet_by_title(src_target_file["sheetname"][0]).get_all_records()
156 # Get transactions file(s) location data
157 df = None
158 transaction_data_file = data_source_router(app_settings_obj.data_source,
159 "transactions",
160 app_settings_obj.g_creds,
161 app_settings_obj.data_sheet, None, None)
162 if transaction_data_file["filetype"] == "csv":
163 # open one or multiple csv files and return a pandas dataframe
164 df = read_csv_as_df(transaction_data_file["filename"])
165 else:
166 # Open one or multiple gsheet worksheets and return a pandas dataframe
167 df = read_gsheet_as_df(transaction_data_file["sheetname"], transaction_data_file["gsheets_obj"])
169 except Exception as e:
170 logger.error(f"Unable to open data source: {app_settings_obj.data_source}. \
171 Please check your names and try again. Error was {e}")
172 raise SystemExit
174 # Clean up value formatting
175 df.reset_index(inplace=True, drop=True)
176 df = df.transform(normalize_amounts, axis=1)
178 is_valid = DataRow.validate(df.columns.to_list(), True)
179 if not is_valid[0]:
180 raise Exception(f"Source data is not in correct format! Message was: {is_valid[1]}")
182 return src_target, df
185def read_csv_as_df(csv_file):
186 if type(csv_file) is list:
187 main_csv = csv_file[0]
188 addl_csvs = csv_file[1:]
189 else:
190 main_csv = csv_file
191 addl_csvs = []
192 if not main_csv.endswith('.csv') or not path.isfile(main_csv):
193 raise Exception(f"Supplied CSV file ({main_csv}) was not found or is the wrong format.")
194 main_df = pd.DataFrame(pd.read_csv(main_csv))
195 for i in addl_csvs:
196 if not i.endswith('.csv') or not path.isfile(i):
197 raise Exception(f"Supplied CSV file ({i}) was not found or is the wrong format.")
198 add_df = pd.DataFrame(pd.read_csv(i))
199 main_df = pd.concat([main_df, add_df], axis=0)
200 return main_df
203def read_gsheet_as_df(gsheet_file, gsheet_obj):
204 # .get_as_df(value_render=pygsheets.ValueRenderOption.UNFORMATTED_VALUE)
205 # # <-- Using unformatted value rounded decimal amounts - not sure why
206 if type(gsheet_file) is list:
207 main_sheet = gsheet_file[0]
208 addl_sheets = gsheet_file[1:]
209 else:
210 main_sheet = gsheet_file
211 addl_sheets = []
212 main_df = gsheet_obj.worksheet_by_title(main_sheet).get_as_df()
213 for i in addl_sheets:
214 add_df = gsheet_obj.worksheet_by_title(i).get_as_df()
215 main_df = pd.concat([main_df, add_df], axis=0)
216 return main_df