bluechips/lib/totals.py

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