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 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 = meta.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 = meta.Session.query(model.Split).\
  34         add_column(sqlalchemy.func.sum(model.Split.share).label('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 = meta.Session.query(model.Transfer).\
  46         add_column(sqlalchemy.func.sum(model.Transfer.amount).label('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 / 100)
  52     for transfer, total_amount in total_credits:
  53         debts[transfer.creditor] += (total_amount / 100)
  54
  55     return debts
  56
  57 def settle(debts_dict):
  58     # This algorithm has been shamelessly stolen from Nelson Elhage's
  59     # <nelhage@mit.edu> implementation for our 2008 summer apartment.
  60
  61     debts_list = [dict(who=user, amount=amount) for user, amount in \
  62                       debts_dict.iteritems()]
  63     debts_list.sort(reverse=True, key=(lambda x: abs(x['amount'])))
  64
  65     owes_list = [debt for debt in debts_list if debt['amount'] > 0]
  66     owed_list = [debt for debt in debts_list if debt['amount'] < 0]
  67
  68     settle = []
  69
  70     while len(owes_list) > 0 and len(owed_list) > 0:
  71         owes = owes_list[0]
  72         owed = owed_list[0]
  73
  74         sum = owes['amount'] + owed['amount']
  75         if sum == 0:
  76             # Perfect balance!
  77             owes_list.pop(0)
  78             owed_list.pop(0)
  79             val = owes['amount']
  80         elif sum > 0:
  81             # person in owes still owes money
  82             owes['amount'] += owed['amount']
  83             owed_list.pop(0)
  84             val = -owed['amount']
  85         else:
  86             # person in owed is owed more than owes has to give
  87             owed['amount'] += owes['amount']
  88             owes_list.pop(0)
  89             val = owes['amount']
  90
  91         settle.append((owes['who'], owed['who'], val))
  92
  93     if len(owes_list) > 0:
  94         raise DirtyBooks, ("People still owe money", owes_list)
  95     if len(owed_list) > 0:
  96         raise DirtyBooks, ("People are still owed money", owed_list)
  97
  98     return settle
  99
 100 __all__ = ['debts', 'settle']