flowrate / flowrate /

import calendar
import datetime
import decimal
import sys

import flowrate
from flowrate.variables import ReferenceFinder, environment as env
from flowrate import cmpdates, ledger

# A table of (, Flow()) pairs.
flows = {}

class Flow(object):

    fields = ('credit', 'debit', 'description',
              'period', 'unit', 'days', 'variables', 'amount',
              'start', 'end')

    def __init__(self, id, **kwargs): = id
        self.row = dict((k, kwargs[k]) for k in self.fields if k != 'variables')

        if 'variables' in kwargs:
            self.row['variables'] = kwargs['variables']
            self.row['variables'] = list(self.get_refs())

        self.obligations = []

    def load_all(cls):
        """Populate the flows dictionary from the database."""
        for row in flowrate.db.execute("SELECT * FROM flows;").fetchall():
            data = {}
            for k in cls.fields:
                f = k
                if f in ('start', 'end'):
                    f = 'range_' + f
                elif f in ('credit', 'debit'):
                    f = f + '_account'
                data[k] = row[f]
            flow = cls(, **data)
            flows[] = flow

    def __getitem__(self, key):
        return self.row[key]

    def __setitem__(self, key, value):
        self.row[key] = value

    def get_refs(self):
        """Return the set of all variable names referenced by this flow."""
        refs = set()
        finder = ReferenceFinder()
        for key in ('start', 'end', 'amount'):
            val = self.row[key]
            if isinstance(val, basestring):
                # Allow 'amount' entries to be numbers
                val = val.strip()
                if val.startswith("="):
                    newrefs = finder.find(val[1:])
        return refs

    def save(self):
        self.row['variables'] = list(self.get_refs())

        if is None:
            row = flowrate.db.execute(
                "INSERT INTO flows"
                " (amount, credit_account, debit_account, range_start, range_end,"
                " period, unit, days, description, variables) "
                "VALUES (%(amount)s, %(credit)s, %(debit)s,"
                " %(start)s, %(end)s, %(period)s, %(unit)s, %(days)s,"
                " %(description)s, %(variables)s) "
                "RETURNING *", self.row).fetchone()
            flows[] = self
            row = {'id':}
                "UPDATE flows SET amount = %(amount)s, "
                "credit_account = %(credit)s, "
                "debit_account = %(debit)s, "
                "range_start = %(start)s, range_end = %(end)s, "
                "period = %(period)s, unit = %(unit)s, days = %(days)s, "
                "description = %(description)s, variables = %(variables)s "
                "WHERE id = %(id)s;", row)

    def delete(self):
        flowrate.db.execute("DELETE FROM flows WHERE id = %s;", (,))
        flows.pop(, None)

    def get_obligations(self):
        """Yield obligation rows for the given flow."""
        credit_type = ledger.accounts[self['credit']]['type']
        credit_mult = -1 if credit_type in ('asset', 'expense') else 1
        debit_type = ledger.accounts[self['debit']]['type']
        debit_mult = 1 if debit_type in ('asset', 'expense') else -1

        unit = cmpdates.units[self['unit']]
        day = cmpdates.days[self['unit']]

        start = self['start'].strip()
        if start.startswith("="):
            # It's a Python expression
            start = env.eval(start[1:], locals())
            # Assume it's an ISO date YYYY-MM-DD
            start =*map(int, start.split("-")))

        end = self['end'].strip()
        if end.startswith("="):
            # It's a Python expression
            end = env.eval(end[1:], locals())
            # Assume it's an ISO date YYYY-MM-DD
            end =*map(int, end.split("-")))

        for d in range(0, (end - start).days + 1):
            postdate = start + datetime.timedelta(days=d)
            if day(postdate) not in self['days']:

            p = unit(postdate)
            # This will be an integer number of units between post and start
            diff = p - unit(start)
            # But start might be after one or more of our
            for fd in sorted(self['days']):
                if day(start) > fd:
                    diff -= 1

            if diff % self['period'] == 0:
                # Eval the amount inside this loop to allow it to refer
                # to the postdate or the iteration variable "d"
                amount = self['amount'].strip()
                if amount.startswith("="):
                    # It's a Python expression
                    amount = env.eval(amount[1:], locals())
                    # Assume it's a number
                    amount = decimal.Decimal(amount)

                # Yield one obligation for the year/month/week on the day.
                yield {
                    'postdate': postdate,
                    'credit': self['credit'],
                    'credit_mult': credit_mult,
                    'debit': self['debit'],
                    'debit_mult': debit_mult,
                    'description': self['description'],
                    'amount': amount,
                    'remaining': amount

    def obligate(self):
        """Insert obligations for self; fulfill as possible."""
        coerce_date = cmpdates.strkeys[self['unit']]

        obs = list(self.get_obligations())

        # Now, fulfill the new obligations.
        # Start by finding the subset of transactions which match accounts;
        # this will be the same for each obligation since they all have the
        # same flow.
        txtable = {}
        for tx in ledger.transactions.itervalues():
            if (flowrate.isSubAccount(tx['credit'], self['credit']) and
                flowrate.isSubAccount(tx['debit'], self['debit'])):
                txtable.setdefault(coerce_date(tx['postdate']), []).append(tx)
        # Order by debit account descending so that
        # more specific accounts match before more general ones.
        for d, txs in txtable.iteritems():
            txs.sort(key=lambda tx: (0 - tx['debit'], tx['postdate']))

        for ob in obs:
            obrem = ob['remaining']
            obdate = coerce_date(ob['postdate'])

            for tx in txtable.get(obdate, []):
                used = sum(f['amount'] for f in tx.fulfillments)
                txrem = tx['amount'] - used
                if obrem > 0 and txrem > 0:
                    f_amt = min(obrem, txrem)
                    tx.fulfillments.append({'obligation': ob, 'amount': f_amt})
                    obrem -= f_amt
                    if obrem <= 0:
                        # Don't allow obligations.remaining to be < 0
                        obrem = 0

            ob['remaining'] = obrem

        # Only swap in the new obligations once all the
        # fulfillment calculations are done.
        self.obligations = obs

    def clear_obligations(self):
        """Remove all existing obligations (and their fulfillments) for self."""
        # TODO: This isn't quite right; it only deletes existing fulfillments
        # for this flow. It might want to delete all fulfillments which
        # *might* apply to this flow, so self.obligate() can "take over"
        # other existing fulfillments. Tough nut.
        for tx in ledger.transactions.itervalues():
            tx.fulfillments = [f for f in tx.fulfillments
                               if f['obligation'] not in self.obligations]
        self.obligations = []

def fulfill(tx):
    """Use the given transaction to fulfill an obligation, if possible.

    Any existing fulfillments for the given transaction will be deleted.
    Then obligations which the transactions might fulfill are fulfilled
    in order until the full amount of the transaction has been applied,
    or no more matching obligations remain.

    Obligations are ordered first by debit account, descending, so that
    more specific expense obligations (like "Restaurants") get filled
    first, but more generic obligations (like "Food") can still be
    fulfilled. It's assumed that this feature is more useful for expenses
    than for income. They are then ordered by date, ascending, within each
    debit account.
    if tx['amount'] <= 0:

    allobs = []
    for flowid, flow in flows.iteritems():
        unit = cmpdates.units[flow['unit']]

        for ob in flow.obligations:
            if (flowrate.isSubAccount(tx['credit'], ob['credit']) and
                flowrate.isSubAccount(tx['debit'], ob['debit']) and
                unit(tx['postdate']) == unit(ob['postdate'])

    txrem = tx['amount']
    for ob in flowrate.ordered(allobs, '-debit', 'postdate'):
        f_amt = min(ob['remaining'], txrem)
        if f_amt > 0:
            tx.fulfillments.append({'obligation': ob, 'amount': f_amt})
            ob['remaining'] -= f_amt
            txrem -= f_amt
            if txrem <= 0:

def unfulfill(tx):
    """Delete any fulfillments for the given transaction. Update obligations."""
    while tx.fulfillments:
        f = tx.fulfillments.pop()
        f['obligation']['remaining'] += f['amount']

def find_transactions(filters):
    """Return flow transactions matching the given criteria (by postdate desc)."""
    if not filters:
        # If there are no filters, return nothing instead of everything
        return []

    # Don't let old unfulfilled obligations screw up balances
    txs = []
    age = - datetime.timedelta(days=30)
    for flowid, flow in flows.iteritems():
        for ob in flow.obligations:
            if ob['remaining'] > 0 and ob['postdate'] > age:
                if all(f(ob) for f in filters):
                        'id': None,
                        'postdate': ob['postdate'],
                        'credit': ob['credit'],
                        'debit': ob['debit'],
                        'description': ob['description'],
                        'amount': ob['remaining'],
    return flowrate.ordered(txs, '-postdate')