Source

flowrate / flowrate / flows.py

Full commit
import calendar
import datetime
import decimal

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


class month(object):
    """A given month in a given year. Like datetime.date without a .day field."""

    __slots__ = ('year', 'month')

    def __init__(self, year, month):
        self.year = year
        self.month = month

    @classmethod
    def fromdate(cls, dt):
        return cls(dt.year, dt.month)

    def todate(self, day=1):
        return datetime.date(self.year, self.month, day)

    def __add__(self, other):
        # Add an integer number of months to self.
        y, m = divmod(self.month + other, 12)
        return month(self.year + y, m)

    def __sub__(self, other):
        if isinstance(other, month):
            # Subtract a month() instance to return an integer diff of months
            y, m = divmod(self.month - other.month, 12)
            return (y * 12) + m
        else:
            # Subtract an integer number of months from self.
            y, m = divmod(self.month - other, 12)
            return month(self.year + y, m)

    def __cmp__(self, other):
        return cmp((self.year, self.month), (other.year, other.month))


class week(object):
    """A given week in a given year. Like datetime.date with a .week field."""

    __slots__ = ('year', 'week')

    def __init__(self, year, week):
        self.year = year
        self.week = week

    @classmethod
    def fromdate(cls, dt):
        y, w, d =  dt.isocalendar()
        return cls(y, w)

    def todate(self, day=1):
        jan1 = datetime.date(self.year, 1, 1)
        day1ofweek1 = jan1 + datetime.timedelta(days=8 - jan1.isocalendar()[2])
        return (day1ofweek1 +
                datetime.timedelta(
                    days=((self.week - 1) * 7) + (day - 1)
                ))

    def __add__(self, other):
        # Add an integer number of weeks to self.
        f = self.todate() + datetime.timedelta(days=other * 7)
        y, w, d = f.isocalendar()
        return week(y, w)

    def __sub__(self, other):
        if isinstance(other, week):
            # Subtract a week() instance to return an integer diff of weeks
            return ((self.todate() - other.todate()).days / 7)
        else:
            # Subtract an integer number of weeks from self.
            f = self.todate() - datetime.timedelta(days=other * 7)
            y, w, d = f.isocalendar()
            return week(y, w)

    def __cmp__(self, other):
        return cmp((self.year, self.week), (other.year, other.week))


class Flow(object):

    def __init__(self, id):
        self.id = id
        self.row = None

    def find(self):
        """Set self.row to a DB row matching the given ID, or None."""
        self.row = flowrate.db.execute(
            "SELECT * FROM flows WHERE id = %s;", (self.id,)).fetchone()

    @classmethod
    def insert(cls, **vals):
        flow = cls(None)

        refs = set()
        finder = ReferenceFinder()
        for key in ('start', 'end', 'amount'):
            val = vals[key].strip()
            if val.startswith("="):
                newrefs = finder.find(val[1:])
                refs.update(newrefs)

        flow.row = flowrate.db.execute("INSERT INTO flows"
                   " (amount, credit_account, debit_account, "
                   "range_start, range_end, period, unit, days, description, "
                   "variables) "
                   "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) "
                   "RETURNING *",
                   (vals["amount"],
                    vals["credit"], vals["debit"],
                    vals["start"], vals["end"],
                    vals["period"], vals['unit'], vals["days"],
                    vals["description"], list(refs),
                    )).fetchone()
        flow.id = flow.row.id
        return flow

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

    def __getattr__(self, key):
        return getattr(self.row, key)

    def obligations(self):
        """Yield obligation rows for the given flow."""
        if self.unit == 'months':
            unit, day = month.fromdate, lambda d: d.day
        elif self.unit == 'weeks':
            unit, day = week.fromdate, lambda d: d.weekday()
        elif self.unit == 'years':
            unit, day = lambda d: d.year, lambda d: d.timetuple().tm_yday

        escrow = datetime.date(2014, 2, 1)

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

        end = self.range_end.strip()
        if end.startswith("="):
            # It's a Python expression
            end = eval(end[1:], env, locals())
        else:
            # Assume it's an ISO date YYYY-MM-DD
            end = datetime.date(*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:
                continue

            p = unit(postdate)
            # This will be an integer number of units between post and start
            diff = p - unit(start)
            # But range_start might be after one or more of our self.day(s)
            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 = eval(amount[1:], env, locals())
                else:
                    # Assume it's a number
                    amount = decimal.Decimal(amount)

                # Yield one transaction for the year/month/week on the day.
                yield {
                    'id': None,
                    'postdate': postdate,
                    'credit': self.credit_account,
                    'debit': self.debit_account,
                    'description': self.description,
                    'amount': amount,
                    }

    def obligate(self):
        """Insert obligation rows for the given flow; fulfill as possible."""
        if self.unit == 'years':
            dategroupformat = 'YYYY'
        elif self.unit == 'days':
            dategroupformat = 'YYYY-MM-DD'
        else:
            dategroupformat = 'YYYY-MM'

        credit_type = flowrate.db.execute(
            "SELECT type FROM accounts WHERE id = %s",
            (self.credit_account,)).fetchone()
        debit_type = flowrate.db.execute(
            "SELECT type FROM accounts WHERE id = %s",
            (self.debit_account,)).fetchone()

        obs = {}
        for ob in self.obligations():
            print ob
            # TODO: this is slow. Can we change it to an "INSERT INTO ... FROM"?
            row = flowrate.db.execute(
                "INSERT INTO obligations "
                "(flowid, postdate, credit_account, debit_account,"
                " description, amount, remaining, dategroupformat,"
                " credit_mult, debit_mult) "
                "VALUES (%s, %s, %s, %s, %s, %s, 0.0, %s, %s, %s) "
                "RETURNING id;",
                (self.id, ob['postdate'], ob['credit'], ob['debit'],
                 ob['description'], ob['amount'], dategroupformat,
                 -1 if credit_type['type'] in ('asset', 'expense') else 1,
                 1 if debit_type['type'] in ('asset', 'expense') else -1,
                 )).fetchone()
            obs[row.id] = ob

        # Now, fulfill the new obligations
        for obid, ob in obs.iteritems():
            obrem = ob['amount']
            for tx in flowrate.db.execute(
                "SELECT t.*, "
                "(SELECT COALESCE(SUM(f.amount), 0) FROM fulfillments f"
                " WHERE f.transactionid = t.id) AS fulfilled "
                "FROM transactions t "
                "WHERE isSubAccount(t.credit_account, %s) "
                "AND isSubAccount(t.debit_account, %s) "
                "AND (to_char(t.postdate, %s) = to_char(%s, %s)) "
                "ORDER BY t.debit_account DESC, t.postdate ASC;",
                (ob['credit'], ob['debit'],
                 dategroupformat, ob['postdate'], dategroupformat)).fetchall(
                ):
                txrem = tx.amount - tx.fulfilled
                if obrem > 0 and txrem > 0:
                    f_amt = min(obrem, txrem)
                    flowrate.db.execute(
                        "INSERT INTO fulfillments "
                        "(transactionid, obligationid, amount) "
                        "VALUES (%s, %s, %s);", (tx.id, obid, f_amt))
                    obrem -= f_amt
                    if obrem <= 0:
                        # Don't allow obligations.remaining to be < 0, below
                        obrem = 0
                        break

            flowrate.db.execute(
                "UPDATE obligations SET remaining = %s "
                "WHERE id = %s;", (obrem, obid))

    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 it can "take over" other existing
        # fulfillments. Tough nut.
        flowrate.db.execute(
            "DELETE FROM fulfillments "
            "WHERE obligationid"
            " IN (SELECT id FROM obligations WHERE flowid = %s); "

            "DELETE FROM obligations WHERE flowid = %s;",
            (self.id, self.id))


def isSubAccount(child, parents):
    for p in parents:
        for scale in (1000, 100, 10, 1):
            if p % scale == 0 and p <= child < (p + scale):
                return True
    return False

def fulfill(txrow):
    """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 txrow.amount <= 0:
        return

    obs = [(row.id, row.amount, row.remaining)
           for row in flowrate.db.execute(
                "SELECT * FROM obligations "
                "WHERE isSubAccount(%s, credit_account) "
                "AND isSubAccount(%s, debit_account) "
                "AND (to_char(%s, dategroupformat) = "
                     "to_char(postdate, dategroupformat)) "
                "ORDER BY debit_account DESC, postdate ASC",
                (txrow.credit_account, txrow.debit_account,
                 txrow.postdate)).fetchall()]

    txrem = txrow.amount
    for obid, obamount, obrem in obs:
        f_amt = min(obrem, txrem)
        if f_amt > 0:
            flowrate.db.execute(
                "INSERT INTO fulfillments "
                "(transactionid, obligationid, amount) VALUES (%s, %s, %s);",
                (txrow.id, obid, f_amt))
            flowrate.db.execute(
                "UPDATE obligations SET remaining = remaining - %s "
                "WHERE id = %s;", (f_amt, obid))
            txrem -= f_amt
            if txrem <= 0:
                break

def unfulfill(txid):
    """Delete any fulfillments for the given transaction. Update obligations."""
    obids = [row.obligationid for row in flowrate.db.execute(
        "DELETE FROM fulfillments "
        "WHERE transactionid = %s RETURNING obligationid;",
        (txid,)).fetchall()]
    for obid in obids:
        flowrate.db.execute(
            "UPDATE obligations ob SET ob.remaining = ob.amount - "
            "(SELECT COALESCE(SUM(f.amount), 0) FROM fulfillments f"
            " WHERE f.obligationid = ob.id) "
            "WHERE ob.id = %s;", (obid,))


def transactions(accounts=None, credits=None, debits=None,
                 years=None, months=None, days=None,
                 description=None):
    """Yield flow transactions matching the given criteria (by postdate desc)."""
    whereclause, args = [], {}
    if accounts:
        whereclause.append("(ARRAY[credit_account] <@ %(accounts)s"
                           " OR ARRAY[debit_account] <@ %(accounts)s)")
        all_accts = [row.id for row in flowrate.db.execute(
                     "SELECT id FROM accounts;").fetchall()]
        args['accounts'] = [a for a in all_accts if isSubAccount(a, accounts)]
    if credits:
        whereclause.append("ARRAY[credit_account] <@ %(credits)s")
        args['credits'] = credits
    if debits:
        whereclause.append("ARRAY[debit_account] <@ %(debits)s")
        args['debits'] = debits
    if description:
        whereclause.append("description ILIKE %(desc)s")
        args['desc'] = '%' + description + '%';
    if years:
        whereclause.append("ARRAY[EXTRACT(year FROM postdate)::integer] <@ %(years)s")
        args['years'] = years
    if months:
        whereclause.append("ARRAY[EXTRACT(month FROM postdate)::integer] <@ %(months)s")
        args['months'] = months
    if days:
        whereclause.append("ARRAY[EXTRACT(day FROM postdate)::integer] <@ %(days)s")
        args['days'] = days

    if not whereclause:
        # If there are no filters, return nothing instead of everything
        return

    for ob in flowrate.db.execute(
        "SELECT * FROM obligations "
        "WHERE " + " AND ".join(whereclause) +
        # Don't let old unfulfilled obligations screw up balances
        " AND postdate >= (CURRENT_DATE - '1 month'::interval)"
        " AND remaining > 0 "
        "ORDER BY postdate DESC;", args).fetchall(
        ):
        yield {
            'postdate': ob.postdate,
            'credit': ob.credit_account,
            'debit': ob.debit_account,
            'description': ob.description,
            'amount': ob.remaining,
            }