1"""Portfolio list extension for Fava.
2
3This is a simple example of Fava's extension reports system.
4"""
5import re
6
7from beancount.core.data import Open
8from beancount.core.number import Decimal
9from beancount.core.number import ZERO
10
11from fava.ext import FavaExtensionBase
12from fava.helpers import FavaAPIException
13from fava.template_filters import cost_or_value
14
15
16class PortfolioList(FavaExtensionBase):  # pragma: no cover
17    """Sample Extension Report that just prints out an Portfolio List."""
18
19    report_title = "Portfolio List"
20
21    def portfolio_accounts(self):
22        """An account tree based on matching regex patterns."""
23        tree = self.ledger.root_tree
24        portfolios = []
25
26        for option in self.config:
27            opt_key = option[0]
28            if opt_key == "account_name_pattern":
29                portfolio = self._account_name_pattern(tree, option[1])
30            elif opt_key == "account_open_metadata_pattern":
31                portfolio = self._account_metadata_pattern(
32                    tree, option[1][0], option[1][1]
33                )
34            else:
35                raise FavaAPIException("Portfolio List: Invalid option.")
36            portfolios.append(portfolio)
37
38        return portfolios
39
40    def _account_name_pattern(self, tree, pattern):
41        """
42        Returns portfolio info based on matching account name.
43
44        Args:
45            tree: Ledger root tree node.
46            pattern: Account name regex pattern.
47        Return:
48            Data structured for use with a querytable (types, rows).
49        """
50        title = "Account names matching: '" + pattern + "'"
51        selected_accounts = []
52        regexer = re.compile(pattern)
53        for acct in tree.keys():
54            if (regexer.match(acct) is not None) and (
55                acct not in selected_accounts
56            ):
57                selected_accounts.append(acct)
58
59        selected_nodes = [tree[x] for x in selected_accounts]
60        portfolio_data = self._portfolio_data(selected_nodes)
61        return title, portfolio_data
62
63    def _account_metadata_pattern(self, tree, metadata_key, pattern):
64        """
65        Returns portfolio info based on matching account open metadata.
66
67        Args:
68            tree: Ledger root tree node.
69            metadata_key: Metadata key to match for in account open.
70            pattern: Metadata value's regex pattern to match for.
71        Return:
72            Data structured for use with a querytable - (types, rows).
73        """
74        title = (
75            "Accounts with '"
76            + metadata_key
77            + "' metadata matching: '"
78            + pattern
79            + "'"
80        )
81        selected_accounts = []
82        regexer = re.compile(pattern)
83        for entry in self.ledger.all_entries_by_type[Open]:
84            if (metadata_key in entry.meta) and (
85                regexer.match(entry.meta[metadata_key]) is not None
86            ):
87                selected_accounts.append(entry.account)
88
89        selected_nodes = [tree[x] for x in selected_accounts]
90        portfolio_data = self._portfolio_data(selected_nodes)
91        return title, portfolio_data
92
93    def _portfolio_data(self, nodes):
94        """
95        Turn a portfolio of tree nodes into querytable-style data.
96
97        Args:
98            nodes: Account tree nodes.
99        Return:
100            types: Tuples of column names and types as strings.
101            rows: Dictionaries of row data by column names.
102        """
103        operating_currency = self.ledger.options["operating_currency"][0]
104        acct_type = ("account", str(str))
105        bal_type = ("balance", str(Decimal))
106        alloc_type = ("allocation", str(Decimal))
107        types = [acct_type, bal_type, alloc_type]
108
109        rows = []
110        portfolio_total = ZERO
111        for node in nodes:
112            row = {}
113            row["account"] = node.name
114            balance = cost_or_value(node.balance)
115            if operating_currency in balance:
116                balance_dec = balance[operating_currency]
117                portfolio_total += balance_dec
118                row["balance"] = balance_dec
119            rows.append(row)
120
121        for row in rows:
122            if "balance" in row:
123                row["allocation"] = round(
124                    (row["balance"] / portfolio_total) * 100, 2
125                )
126
127        return types, rows
128