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 = {}
  25
  26     # First, credit everyone for expenditures they've made
  27     for user in users:
  28         debts_dict[user] = Currency(-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_dict[split.user] += total_cents
  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_dict[transfer.debtor] -= total_amount
  52     for transfer, total_amount in total_credits:
  53         debts_dict[transfer.creditor] += total_amount
  54
  55     return debts_dict
  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_list = []
  69
  70     while len(owes_list) > 0 and len(owed_list) > 0:
  71         owes_list.sort(reverse=True, key=(lambda x: abs(x['amount'])))
  72         owed_list.sort(reverse=True, key=(lambda x: abs(x['amount'])))
  73
  74         owes = owes_list[0]
  75         owed = owed_list[0]
  76
  77         sum = owes['amount'] + owed['amount']
  78         if sum == 0:
  79             # Perfect balance!
  80             owes_list.pop(0)
  81             owed_list.pop(0)
  82             val = owes['amount']
  83         elif sum > 0:
  84             # person in owes still owes money
  85             owes['amount'] += owed['amount']
  86             owed_list.pop(0)
  87             val = -owed['amount']
  88         else:
  89             # person in owed is owed more than owes has to give
  90             owed['amount'] += owes['amount']
  91             owes_list.pop(0)
  92             val = owes['amount']
  93
  94         settle_list.append((owes['who'], owed['who'], val))
  95
  96     if len(owes_list) > 0:
  97         raise DirtyBooks, ("People still owe money", owes_list)
  98     if len(owed_list) > 0:
  99         raise DirtyBooks, ("People are still owed money", owed_list)
 100
 101     return settle_list
 102
 103 __all__ = ['debts', 'settle']