1"""Getter functions that operate on lists of entries to return various lists of
2things that they reference, accounts, tags, links, currencies, etc.
3"""
4__copyright__ = "Copyright (C) 2013-2016  Martin Blais"
5__license__ = "GNU GPLv2"
6
7from collections import defaultdict
8from collections import OrderedDict
9
10from beancount.core.data import Transaction
11from beancount.core.data import Open
12from beancount.core.data import Close
13from beancount.core.data import Commodity
14from beancount.core import account
15
16
17class GetAccounts:
18    """Accounts gatherer.
19    """
20    def get_accounts_use_map(self, entries):
21        """Gather the list of accounts from the list of entries.
22
23        Args:
24          entries: A list of directive instances.
25        Returns:
26          A pair of dictionaries of account name to date, one for first date
27          used and one for last date used. The keys should be identical.
28        """
29        accounts_first = {}
30        accounts_last = {}
31        for entry in entries:
32            method = getattr(self, entry.__class__.__name__)
33            for account_ in method(entry):
34                if account_ not in accounts_first:
35                    accounts_first[account_] = entry.date
36                accounts_last[account_] = entry.date
37        return accounts_first, accounts_last
38
39    def get_entry_accounts(self, entry):
40        """Gather all the accounts references by a single directive.
41
42        Note: This should get replaced by a method on each directive eventually,
43        that would be the clean way to do this.
44
45        Args:
46          entry: A directive instance.
47        Returns:
48          A set of account name strings.
49        """
50        method = getattr(self, entry.__class__.__name__)
51        return set(method(entry))
52
53    # pylint: disable=invalid-name
54
55    def Transaction(_, entry):
56        """Process a Transaction directive.
57
58        Args:
59          entry: An instance of Transaction.
60        Yields:
61          The accounts of the legs of the transaction.
62        """
63        for posting in entry.postings:
64            yield posting.account
65
66    def Pad(_, entry):
67        """Process a Pad directive.
68
69        Args:
70          entry: An instance of Pad.
71        Returns:
72          The two accounts of the Pad directive.
73        """
74        return (entry.account, entry.source_account)
75
76    def _one(_, entry):
77        """Process directives with a single account attribute.
78
79        Args:
80          entry: An instance of a directive.
81        Returns:
82          The single account of this directive.
83        """
84        return (entry.account,)
85
86    def _zero(_, entry):
87        """Process directives with no accounts.
88
89        Args:
90          entry: An instance of a directive.
91        Returns:
92          An empty list
93        """
94        return ()
95
96    # Associate all the possible directives with their respective handlers.
97    Open = Close = Balance = Note = Document = _one
98    Commodity = Event = Query = Price = Custom = _zero
99
100
101# Global instance to share.
102_GetAccounts = GetAccounts()
103
104
105def get_accounts_use_map(entries):
106    """Gather all the accounts references by a list of directives.
107
108    Args:
109      entries: A list of directive instances.
110    Returns:
111      A pair of dictionaries of account name to date, one for first date
112      used and one for last date used. The keys should be identical.
113    """
114    return _GetAccounts.get_accounts_use_map(entries)
115
116
117def get_accounts(entries):
118    """Gather all the accounts references by a list of directives.
119
120    Args:
121      entries: A list of directive instances.
122    Returns:
123      A set of account strings.
124    """
125    _, accounts_last = _GetAccounts.get_accounts_use_map(entries)
126    return accounts_last.keys()
127
128
129def get_entry_accounts(entry):
130    """Gather all the accounts references by a single directive.
131
132    Note: This should get replaced by a method on each directive eventually,
133    that would be the clean way to do this.
134
135    Args:
136      entries: A directive instance.
137    Returns:
138      A set of account strings.
139    """
140    return _GetAccounts.get_entry_accounts(entry)
141
142
143def get_account_components(entries):
144    """Gather all the account components available in the given directives.
145
146    Args:
147      entries: A list of directive instances.
148    Returns:
149      A list of strings, the unique account components, including the root
150      account names.
151    """
152    accounts = get_accounts(entries)
153    components = set()
154    for account_name in accounts:
155        components.update(account.split(account_name))
156    return sorted(components)
157
158
159def get_all_tags(entries):
160    """Return a list of all the tags seen in the given entries.
161
162    Args:
163      entries: A list of directive instances.
164    Returns:
165      A set of tag strings.
166    """
167    all_tags = set()
168    for entry in entries:
169        if not isinstance(entry, Transaction):
170            continue
171        if entry.tags:
172            all_tags.update(entry.tags)
173    return sorted(all_tags)
174
175
176def get_all_payees(entries):
177    """Return a list of all the unique payees seen in the given entries.
178
179    Args:
180      entries: A list of directive instances.
181    Returns:
182      A set of payee strings.
183    """
184    all_payees = set()
185    for entry in entries:
186        if not isinstance(entry, Transaction):
187            continue
188        all_payees.add(entry.payee)
189    all_payees.discard(None)
190    return sorted(all_payees)
191
192
193def get_all_links(entries):
194    """Return a list of all the links seen in the given entries.
195
196    Args:
197      entries: A list of directive instances.
198    Returns:
199      A set of links strings.
200    """
201    all_links = set()
202    for entry in entries:
203        if not isinstance(entry, Transaction):
204            continue
205        if entry.links:
206            all_links.update(entry.links)
207    return sorted(all_links)
208
209
210def get_leveln_parent_accounts(account_names, level, nrepeats=0):
211    """Return a list of all the unique leaf names at level N in an account hierarchy.
212
213    Args:
214      account_names: A list of account names (strings)
215      level: The level to cross-cut. 0 is for root accounts.
216      nrepeats: A minimum number of times a leaf is required to be present in the
217        the list of unique account names in order to be returned by this function.
218    Returns:
219      A list of leaf node names.
220    """
221    leveldict = defaultdict(int)
222    for account_name in set(account_names):
223        components = account.split(account_name)
224        if level < len(components):
225            leveldict[components[level]] += 1
226    levels = {level_
227              for level_, count in leveldict.items()
228              if count > nrepeats}
229    return sorted(levels)
230
231
232def get_dict_accounts(account_names):
233    """Return a nested dict of all the unique leaf names.
234    account names are labelled with LABEL=True
235
236    Args:
237      account_names: An iterable of account names (strings)
238    Returns:
239      A nested OrderedDict of account leafs
240    """
241    leveldict = OrderedDict()
242    for account_name in account_names:
243        nested_dict = leveldict
244        for component in account.split(account_name):
245            nested_dict = nested_dict.setdefault(component, OrderedDict())
246        nested_dict[get_dict_accounts.ACCOUNT_LABEL] = True
247    return leveldict
248get_dict_accounts.ACCOUNT_LABEL = '__root__'
249
250
251def get_min_max_dates(entries, types=None):
252    """Return the minimum and maximum dates in the list of entries.
253
254    Args:
255      entries: A list of directive instances.
256      types: An optional tuple of types to restrict the entries to.
257    Returns:
258      A pair of datetime.date dates, the minimum and maximum dates seen in the
259      directives.
260    """
261    date_first = date_last = None
262
263    for entry in entries:
264        if types and not isinstance(entry, types):
265            continue
266        date_first = entry.date
267        break
268
269    for entry in reversed(entries):
270        if types and not isinstance(entry, types):
271            continue
272        date_last = entry.date
273        break
274
275    return (date_first, date_last)
276
277
278def get_active_years(entries):
279    """Yield all the years that have at least one entry in them.
280
281    Args:
282      entries: A list of directive instances.
283    Yields:
284      Unique dates see in the list of directives.
285    """
286    seen = set()
287    prev_year = None
288    for entry in entries:
289        year = entry.date.year
290        if year != prev_year:
291            prev_year = year
292            assert year not in seen
293            seen.add(year)
294            yield year
295
296
297def get_account_open_close(entries):
298    """Fetch the open/close entries for each of the accounts.
299
300    If an open or close entry happens to be duplicated, accept the earliest
301    entry (chronologically).
302
303    Args:
304      entries: A list of directive instances.
305    Returns:
306      A map of account name strings to pairs of (open-directive, close-directive)
307      tuples.
308    """
309    # A dict of account name to (open-entry, close-entry).
310    open_close_map = defaultdict(lambda: [None, None])
311    for entry in entries:
312        if not isinstance(entry, (Open, Close)):
313            continue
314        open_close = open_close_map[entry.account]
315        index = 0 if isinstance(entry, Open) else 1
316        previous_entry = open_close[index]
317        if previous_entry is not None:
318            if previous_entry.date <= entry.date:
319                entry = previous_entry
320        open_close[index] = entry
321
322    return dict(open_close_map)
323
324
325def get_commodity_directives(entries):
326    """Create map of commodity names to Commodity entries.
327
328    Args:
329      entries: A list of directive instances.
330    Returns:
331      A map of commodity name strings to Commodity directives.
332    """
333    return {entry.currency: entry for entry in entries if isinstance(entry, Commodity)}
334
335
336def get_values_meta(name_to_entries_map, *meta_keys, default=None):
337    """Get a map of the metadata from a map of entries values.
338
339    Given a dict of some key to a directive instance (or None), return a mapping
340    of the key to the metadata extracted from each directive, or a default
341    value. This can be used to gather a particular piece of metadata from an
342    accounts map or a commodities map.
343
344    Args:
345      name_to_entries_map: A dict of something to an entry or None.
346      meta_keys: A list of strings, the keys to fetch from the metadata.
347      default: The default value to use if the metadata is not available or if
348        the value/entry is None.
349    Returns:
350      A mapping of the keys of name_to_entries_map to the values of the 'meta_keys'
351      metadata. If there are multiple 'meta_keys', each value is a tuple of them.
352      On the other hand, if there is only a single one, the value itself is returned.
353    """
354    value_map = {}
355    for key, entry in name_to_entries_map.items():
356        value_list = []
357        for meta_key in meta_keys:
358            value_list.append(entry.meta.get(meta_key, default)
359                              if entry is not None
360                              else default)
361        value_map[key] = (value_list[0]
362                          if len(meta_keys) == 1
363                          else tuple(value_list))
364    return value_map
365