1use crate::account_balance::AccountBalance;
2use crate::{Amount, Ledger, Transaction};
3use std::collections::HashMap;
4use std::ops::AddAssign;
5use std::ops::SubAssign;
6
7#[derive(Debug, Clone)]
11pub struct Balance {
12 pub account_balances: HashMap<String, AccountBalance>,
13}
14
15impl Default for Balance {
16 fn default() -> Self {
17 Self::new()
18 }
19}
20
21impl Balance {
22 pub fn new() -> Balance {
23 Balance {
24 account_balances: HashMap::new(),
25 }
26 }
27
28 pub fn update_with_transaction(&mut self, transaction: &Transaction) {
29 for posting in &transaction.postings {
30 let account_balance = self
31 .account_balances
32 .entry(posting.account.clone())
33 .or_default();
34
35 account_balance
36 .amounts
37 .entry(posting.amount.commodity.name.clone())
38 .and_modify(|a| a.quantity += posting.amount.quantity)
39 .or_insert_with(|| posting.amount.clone());
40 }
41 self.remove_empties();
42 }
43
44 pub fn get_account_balance(&self, account_prefixes: &[&str]) -> AccountBalance {
45 let mut balance = AccountBalance::new();
46 for (account_name, account_balance) in &self.account_balances {
47 for account_prefix in account_prefixes {
48 if account_name.starts_with(account_prefix) {
49 balance += account_balance;
50 break;
51 }
52 }
53 }
54
55 balance
56 }
57
58 pub fn add_amount(&mut self, account: &str, amount: &Amount) {
59 let account_balance = self.account_balances.entry(account.to_owned()).or_default();
60 *account_balance += amount;
61 }
62
63 fn remove_empties(&mut self) {
64 let empties: Vec<String> = self
65 .account_balances
66 .iter()
67 .filter(|&(_, account_balance)| account_balance.is_zero())
68 .map(|(k, _)| k.clone())
69 .collect();
70 for empty in empties {
71 self.account_balances.remove(&empty);
72 }
73 }
74}
75
76impl<'a> From<&'a Ledger> for Balance {
77 fn from(ledger: &'a Ledger) -> Self {
78 let mut balance = Balance::new();
79
80 for transaction in &ledger.transactions {
81 balance.update_with_transaction(transaction);
82 }
83
84 balance
85 }
86}
87
88impl<'a> From<&'a Transaction> for Balance {
89 fn from(transaction: &'a Transaction) -> Self {
90 let mut balance = Balance::new();
91 balance.update_with_transaction(transaction);
92 balance
93 }
94}
95
96impl<'a> AddAssign<&'a Balance> for Balance {
97 fn add_assign(&mut self, other: &'a Balance) {
98 for (account_name, account_balance) in &other.account_balances {
99 self.account_balances
100 .entry(account_name.clone())
101 .and_modify(|b| *b += account_balance)
102 .or_insert_with(|| account_balance.clone());
103 }
104 self.remove_empties();
105 }
106}
107
108impl<'a> SubAssign<&'a Balance> for Balance {
109 fn sub_assign(&mut self, other: &'a Balance) {
110 for (account_name, account_balance) in &other.account_balances {
111 self.account_balances
112 .entry(account_name.clone())
113 .and_modify(|b| *b -= account_balance)
114 .or_insert_with(|| account_balance.clone());
115 }
116 self.remove_empties();
117 }
118}