bluechips/lib/totals.py

   1 """
   2 Calculate the total state of the books
   3 """
   4
   5 from bluechips import model
   6 from bluechips.model import meta
   7
   8 from bluechips.model.types import Currency
   9
  10 import sqlalchemy
  11
  12 class DirtyBooks(Exception):
  13     """
  14     If the books don't work out, raise this
  15     """
  16     pass
  17
  18 def debts():
  19     # In this scheme, negative numbers represent money the house owes
  20     # the user, and positive numbers represent money the user owes the
  21     # house
  22     users = meta.Session.query(model.User)
  23
  24     debts_dict = dict((u, Currency(0)) for u in users)
  25
  26     # First, credit everyone for expenditures they've made
  27     total_expenditures = meta.Session.query(model.Expenditure).\
  28         add_column(sqlalchemy.func.sum(model.Expenditure.amount).label('total_spend')).\
  29         group_by(model.Expenditure.spender_id)
  30     for expenditure, total_spend in total_expenditures:
  31         debts_dict[expenditure.spender] -= total_spend
  32
  33     # Next, debit everyone for expenditures that they have an
  34     # investment in (i.e. splits)
  35
  36     total_splits = meta.Session.query(model.Split).\
  37         add_column(sqlalchemy.func.sum(model.Split.share).label('total_split')).\
  38         group_by(model.Split.user_id)
  39
  40     for split, total_cents in total_splits:
  41         debts_dict[split.user] += total_cents
  42
  43     # Finally, move transfers around appropriately
  44     #
  45     # To keep this from getting to be expensive, have SQL sum up
  46     # transfers for us
  47
  48     transfer_q = meta.Session.query(model.Transfer).\
  49         add_column(sqlalchemy.func.sum(model.Transfer.amount).label('total_amount'))
  50     total_debits = transfer_q.group_by(model.Transfer.debtor_id)
  51     total_credits = transfer_q.group_by(model.Transfer.creditor_id)
  52
  53     for transfer, total_amount in total_debits:
  54         debts_dict[transfer.debtor] -= total_amount
  55     for transfer, total_amount in total_credits:
  56         debts_dict[transfer.creditor] += total_amount
  57
  58     return debts_dict
  59
  60 def settle(debts_dict):
  61     # This algorithm has been shamelessly stolen from Nelson Elhage's
  62     # <nelhage@mit.edu> implementation for our 2008 summer apartment.
  63
  64     debts_list = [dict(who=user, amount=amount) for user, amount in \
  65                       debts_dict.iteritems()]
  66     #debts_list.sort(reverse=True, key=(lambda x: abs(x['amount'])))
  67
  68     owes_list = [debt for debt in debts_list if debt['amount'] > 0]
  69     owed_list = [debt for debt in debts_list if debt['amount'] < 0]
  70
  71     settle_list = []
  72
  73     while len(owes_list) > 0 and len(owed_list) > 0:
  74         owes_list.sort(reverse=True, key=(lambda x: abs(x['amount'])))
  75         owed_list.sort(reverse=True, key=(lambda x: abs(x['amount'])))
  76
  77         owes = owes_list[0]
  78         owed = owed_list[0]
  79
  80         sum = owes['amount'] + owed['amount']
  81         if sum == 0:
  82             # Perfect balance!
  83             owes_list.pop(0)
  84             owed_list.pop(0)
  85             val = owes['amount']
  86         elif sum > 0:
  87             # person in owes still owes money
  88             owes['amount'] += owed['amount']
  89             owed_list.pop(0)
  90             val = -owed['amount']
  91         else:
  92             # person in owed is owed more than owes has to give
  93             owed['amount'] += owes['amount']
  94             owes_list.pop(0)
  95             val = owes['amount']
  96
  97         settle_list.append((owes['who'], owed['who'], val))
  98
  99     if len(owes_list) > 0:
 100         raise DirtyBooks, ("People still owe money", owes_list)
 101     if len(owed_list) > 0:
 102         raise DirtyBooks, ("People are still owed money", owed_list)
 103
 104     return settle_list
 105
 106 __all__ = ['debts', 'settle']