ledger_utils/
balance.rs

1use crate::account_balance::AccountBalance;
2use crate::{Amount, Ledger, Transaction};
3use std::collections::HashMap;
4use std::ops::AddAssign;
5use std::ops::SubAssign;
6
7/// Balance of one or more accounts.
8///
9/// Maps account names to their balances.
10#[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}