-
Notifications
You must be signed in to change notification settings - Fork 0
/
statement_puller.py
165 lines (136 loc) · 6.04 KB
/
statement_puller.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import datetime
from pathlib import Path
from typing import List, Tuple
import click
from ofxclient.account import Account
from ofxclient.config import OfxConfig
from category import Category
from category_lookup_table import CategoryLookupTable, CategoryHinter
from transaction import Transaction
from statement_parser import parse
def _prompt_user_for_category(table: CategoryLookupTable, hinter: CategoryHinter, description: str, description_suffix: str) -> Category:
"""
Prompts the user so the user can decide what category this transaction belongs in.
The user can either:
1. Skip (do nothing, category is Unknown)
2. Select the hint for one of the keywords in the description
The keyword to category mapping will be saved in the hints config
The description to category mapping will be saved in the lookup table
3. Select the category for the entire description
The description to category mapping will be saved in the lookup table
Params:
table : Lookup table to store the category for the description.
hinter : Produces hints for the potential categories based on the description.
description : The transaction's description.
description_suffix : A suffix to append to the description.
Returns:
The category if chosen.
"""
SKIP_KEYWORD = "SKIP"
hints = hinter.hint(description)
# Add an option to skip
options: List[Tuple[Category, str]] = [(Category.Unknown, SKIP_KEYWORD)]
# Add the hints first
# This will save the category for the word
options.extend((c, word) for word, c in hints.items())
# Add the options for all categories last
# This will save the category for the whole description
options.extend(((c, "") for c in list(Category)))
# Format the options
# Add enumerations for each option
format_word = lambda word: f"({word})" if word else ""
option_strs = [f"({i}) {c.name} {format_word(word)}" for i, (c, word) in enumerate(options)]
option_lines = "\n".join(option_strs)
options_str = f"\nDescription: {description} {description_suffix}\nChoose which category best fits the description.\n{option_lines}\n"
# Prompt the user for which category this transaction belongs in
# Keep prompting until successful
while True:
option = input(options_str)
try:
category, word = options[int(option)]
# If skipping, use the unknown category
if word == SKIP_KEYWORD:
return category
# If the selection was for a word, store it as a word hint
if word:
hinter.store(word, category)
# Always save the category for this description
table.store(description, category)
return category
except (ValueError, KeyError):
pass
def determine_transaction_categories(
table: CategoryLookupTable,
hinter: CategoryHinter,
transactions: List[Transaction],
*,
do_prompt: bool,
) -> List[Transaction]:
"""
Determines the categories for all unknown categories.
Args:
table : Lookup table to store the category for the description.
hinter : Produces hints for the potential categories based on the description.
transactions : All the transactions to determine categories for.
do_prompt : True to prompt the user for which category is best.
Returns:
The modified transactions.
"""
for i, transaction in enumerate(transactions):
# Try to load from the lookup table first
if category := table.load(transaction.description):
transactions[i].category = category
continue
# Otherwise prompt the user for what the category should be
if do_prompt:
transactions[i].category = _prompt_user_for_category(table, hinter, transaction.description, description_suffix=f"({i} / {len(transactions)})")
return transactions
def _download(account: Account) -> Path:
date = datetime.datetime.now().date().strftime("%Y-%m-%d")
path = Path(__file__).parent / Path("data", date + ".ofx")
if path.exists():
print(f"Found cached statement {path}")
return path
print(f"Downloading statement to {path}")
data = account.download(days=365)
with path.open("wb", encoding="utf-8") as f:
f.write(data.getvalue().encode())
return path
def pull(username: str, account_name: str) -> List[Transaction]:
"""
Pulls the OFX transactions from the specified account.
The [ofxclient] config must be initialized already by running it from the command line.
Params:
username : The name of the account's owner.
account_name : The name of the account to pull from.
Returns:
The transaction if the account was found.
"""
# Load the config and the accounts that belong to the config
config = OfxConfig()
accounts = list(
filter(
lambda a: (a.institution.username == username) and (a.description == account_name),
config.accounts(),
)
)
if not accounts:
raise KeyError(f"Did not find the {account_name} account under user {username}")
if len(accounts) > 1:
raise RuntimeError(f"Found {len(accounts)} accounts for {account_name} under user {username}")
# Try to download as much as possible, but Chase for example only returns 30 days
account = accounts[0]
statement_path = _download(account)
# Parse the transactions, this requires reading from the file in binary mode
transactions = parse(fname=statement_path)
determine_transaction_categories(transactions)
return transactions
@click.command()
@click.argument("username", type=str)
@click.argument("account_name", type=str)
def cli(username: str, account_name: str) -> None:
"""CLI for testing the [pull] function."""
import pprint
pprint.pprint(pull(username=username, account_name=account_name))
if __name__ == "__main__":
cli() # pylint: disable=no-value-for-parameter