/beancount/ops/pad.py

https://github.com/beancount/beancount · Python · 171 lines · 73 code · 18 blank · 80 comment · 19 complexity · adcf9f1565ea473abddcf23c8929d744 MD5 · raw file

  1. """Automatic padding of gaps between entries.
  2. """
  3. __copyright__ = "Copyright (C) 2013-2016 Martin Blais"
  4. __license__ = "GNU GPLv2"
  5. import collections
  6. from beancount.core import account
  7. from beancount.core import amount
  8. from beancount.core import inventory
  9. from beancount.core import data
  10. from beancount.core import position
  11. from beancount.core import flags
  12. from beancount.core import realization
  13. from beancount.utils import misc_utils
  14. from beancount.ops import balance
  15. __plugins__ = ('pad',)
  16. PadError = collections.namedtuple('PadError', 'source message entry')
  17. def pad(entries, options_map):
  18. """Insert transaction entries for to fulfill a subsequent balance check.
  19. Synthesize and insert Transaction entries right after Pad entries in order
  20. to fulfill checks in the padded accounts. Returns a new list of entries.
  21. Note that this doesn't pad across parent-child relationships, it is a very
  22. simple kind of pad. (I have found this to be sufficient in practice, and
  23. simpler to implement and understand.)
  24. Furthermore, this pads for a single currency only, that is, balance checks
  25. are specified only for one currency at a time, and pads will only be
  26. inserted for those currencies.
  27. Args:
  28. entries: A list of directives.
  29. options_map: A parser options dict.
  30. Returns:
  31. A new list of directives, with Pad entries inserted, and a list of new
  32. errors produced.
  33. """
  34. pad_errors = []
  35. # Find all the pad entries and group them by account.
  36. pads = list(misc_utils.filter_type(entries, data.Pad))
  37. pad_dict = misc_utils.groupby(lambda x: x.account, pads)
  38. # Partially realize the postings, so we can iterate them by account.
  39. by_account = realization.postings_by_account(entries)
  40. # A dict of pad -> list of entries to be inserted.
  41. new_entries = {id(pad): [] for pad in pads}
  42. # Process each account that has a padding group.
  43. for account_, pad_list in sorted(pad_dict.items()):
  44. # Last encountered / currency active pad entry.
  45. active_pad = None
  46. # Gather all the postings for the account and its children.
  47. postings = []
  48. is_child = account.parent_matcher(account_)
  49. for item_account, item_postings in by_account.items():
  50. if is_child(item_account):
  51. postings.extend(item_postings)
  52. postings.sort(key=data.posting_sortkey)
  53. # A set of currencies already padded so far in this account.
  54. padded_lots = set()
  55. pad_balance = inventory.Inventory()
  56. for entry in postings:
  57. assert not isinstance(entry, data.Posting)
  58. if isinstance(entry, data.TxnPosting):
  59. # This is a transaction; update the running balance for this
  60. # account.
  61. pad_balance.add_position(entry.posting)
  62. elif isinstance(entry, data.Pad):
  63. if entry.account == account_:
  64. # Mark this newly encountered pad as active and allow all lots
  65. # to be padded heretofore.
  66. active_pad = entry
  67. padded_lots = set()
  68. elif isinstance(entry, data.Balance):
  69. check_amount = entry.amount
  70. # Compare the current balance amount to the expected one from
  71. # the check entry. IMPORTANT: You need to understand that this
  72. # does not check a single position, but rather checks that the
  73. # total amount for a particular currency (which itself is
  74. # distinct from the cost).
  75. balance_amount = pad_balance.get_currency_units(check_amount.currency)
  76. diff_amount = amount.sub(balance_amount, check_amount)
  77. # Use the specified tolerance or automatically infer it.
  78. tolerance = balance.get_balance_tolerance(entry, options_map)
  79. if abs(diff_amount.number) > tolerance:
  80. # The check fails; we need to pad.
  81. # Pad only if pad entry is active and we haven't already
  82. # padded that lot since it was last encountered.
  83. if active_pad and (check_amount.currency not in padded_lots):
  84. # Note: we decide that it's an error to try to pad
  85. # positions at cost; we check here that all the existing
  86. # positions with that currency have no cost.
  87. positions = [pos
  88. for pos in pad_balance.get_positions()
  89. if pos.units.currency == check_amount.currency]
  90. for position_ in positions:
  91. if position_.cost is not None:
  92. pad_errors.append(
  93. PadError(entry.meta,
  94. ("Attempt to pad an entry with cost for "
  95. "balance: {}".format(pad_balance)),
  96. active_pad))
  97. # Thus our padding lot is without cost by default.
  98. diff_position = position.Position.from_amounts(
  99. amount.Amount(check_amount.number - balance_amount.number,
  100. check_amount.currency))
  101. # Synthesize a new transaction entry for the difference.
  102. narration = ('(Padding inserted for Balance of {} for '
  103. 'difference {})').format(check_amount, diff_position)
  104. new_entry = data.Transaction(
  105. active_pad.meta.copy(), active_pad.date, flags.FLAG_PADDING,
  106. None, narration, data.EMPTY_SET, data.EMPTY_SET, [])
  107. new_entry.postings.append(
  108. data.Posting(active_pad.account,
  109. diff_position.units, diff_position.cost,
  110. None, None, None))
  111. neg_diff_position = -diff_position
  112. new_entry.postings.append(
  113. data.Posting(active_pad.source_account,
  114. neg_diff_position.units, neg_diff_position.cost,
  115. None, None, None))
  116. # Save it for later insertion after the active pad.
  117. new_entries[id(active_pad)].append(new_entry)
  118. # Fixup the running balance.
  119. pos, _ = pad_balance.add_position(diff_position)
  120. if pos is not None and pos.is_negative_at_cost():
  121. raise ValueError(
  122. "Position held at cost goes negative: {}".format(pos))
  123. # Mark this lot as padded. Further checks should not pad this lot.
  124. padded_lots.add(check_amount.currency)
  125. # Insert the newly created entries right after the pad entries that created them.
  126. padded_entries = []
  127. for entry in entries:
  128. padded_entries.append(entry)
  129. if isinstance(entry, data.Pad):
  130. entry_list = new_entries[id(entry)]
  131. if entry_list:
  132. padded_entries.extend(entry_list)
  133. else:
  134. # Generate errors on unused pad entries.
  135. pad_errors.append(
  136. PadError(entry.meta, "Unused Pad entry", entry))
  137. return padded_entries, pad_errors