Source

flowrate / flowrate / csvutil.py

Full commit
import csv
import datetime
import decimal
import sqlalchemy

import flowrate
from flowrate import flows


class Peekable(object):
    """A file-like object for peeking on a non-seekable file."""

    def __init__(self, rfile):
        self.rfile = rfile
        self.buffer = ""

    def peek(self, size):
        """Buffer data from the file. Equivalent to read(size); seek(-size)."""
        data = self.rfile.read(size)
        self.buffer += data
        return data

    def read(self, size):
        # Read bytes from the buffer, if present
        data = ""
        if self.buffer:
            data = self.buffer[:size]
            self.buffer = self.buffer[size:]

        # Continue reading from the file, if necessary
        if len(data) < size:
            data += self.rfile.read(size - len(data))

        return data

    def readline(self, size=None):
        data = ""
        if self.buffer:
            pos = self.buffer.find("\n")
            if pos == -1:
                # Consume the whole buffer
                if size is None:
                    data = self.buffer
                    self.buffer = ""
                else:
                    data = self.buffer[:size]
                    self.buffer = self.buffer[size:]
            else:
                # Consume part of the buffer
                # Increment pos by 1 for the newline char itself
                pos += 1
                if size is None or size >= pos:
                    data = self.buffer[:pos]
                    self.buffer = self.buffer[pos:]
                else:
                    data = self.buffer[:size]
                    self.buffer = self.buffer[size:]

        # Continue reading from the file, if necessary
        if "\n" not in data:
            if size is None:
                data += self.rfile.readline()
            elif len(data) < size:
                data += self.rfile.readline(size - len(data))

        return data

    def readlines(self, sizehint=0):
        # Shamelessly stolen from StringIO
        total = 0
        lines = []
        line = self.readline()
        while line:
            lines.append(line)
            total += len(line)
            if 0 < sizehint <= total:
                break
            line = self.readline()
        return lines
    
    def close(self):
        self.rfile.close()
    
    def __iter__(self):
        return self
    
    def __next__(self):
        data = self.readline()
        if data == "":
            raise StopIteration
        return data
    
    def next(self):
        data = self.readline()
        if data == "":
            raise StopIteration
        return data


class Importer(object):

    def get_account(self, account_name):
        return flowrate.db.execute("SELECT * FROM accounts WHERE name = %s",
                                   (account_name,)).fetchone()

    def add_tx(self, credit, debit, amount, postdate, description):
        exists = flowrate.db.execute(
            "SELECT * FROM transactions "
            "WHERE credit_account = %s AND debit_account = %s "
            "AND amount = %s AND postdate = %s;",
            (credit['id'], debit['id'], amount, postdate)).fetchone()
        if exists is None:
            newrow = flowrate.db.execute("INSERT INTO transactions "
                "(amount, credit_account, debit_account, postdate,"
                " description, credit_mult, debit_mult) "
                "VALUES (%s, %s, %s, %s, %s, %s, %s) RETURNING *",
                (amount, credit['id'], debit['id'], postdate, description,
                 -1 if credit['type'] in ('asset', 'expense') else 1,
                 1 if debit['type'] in ('asset', 'expense') else -1,
                 )).fetchone()
            flows.fulfill(newrow)
            return newrow.id
        else:
            return None


class MintImporter(Importer):

    def load_mint_csv(self, f):
        """Read the given Mint file and assert its transactions exist in Flowrate."""
        f = Peekable(f)
        dialect = csv.Sniffer().sniff(f.peek(1024))
        reader = csv.DictReader(f, dialect=dialect)

        for tx in reader:
            self.process_tx(tx)

    def process_tx(self, tx):
        """Process the given transaction row from a CSV DictReader."""
        if tx['Transaction Type'] == 'credit':
            credit = self.get_account(tx['Category'])
            debit = self.get_account(tx['Account Name'])
        else:
            credit = self.get_account(tx['Account Name'])
            debit = self.get_account(tx['Category'])
        # Let subclasses override get_account if they
        # want to add new accountids, otherwise pass
        if credit is None or debit is None:
            return

        amount = decimal.Decimal(tx['Amount'])
        month, day, year = map(int, tx['Date'].split('/'))
        postdate = datetime.date(year, month, day)
        self.add_tx(credit, debit, amount, postdate, tx['Description'])


class ImportException(Exception):
    pass


class UnknownAccount(ImportException):

    def __init__(self, account_name):
        self.account_name = account_name


class ErrorGatheringMintImporter(MintImporter):

    def __init__(self):
        self.errors = []

    def process_tx(self, tx):
        try:
            result = MintImporter.process_tx(self, tx)
        except ImportException, exc:
            exc.tx = tx
            self.errors.append(exc)

    def get_account(self, account_name):
        acct = Importer.get_account(self, account_name)
        if acct is None:
            raise UnknownAccount(account_name)
        return acct


class CommandLineMintImporter(MintImporter):

    def get_account(self, account_name):
        acct = Importer.get_account(self, account_name)
        if acct is None:
            acct = self.ask_account(account_name)
        return acct

    def ask_account(self, account_name):
        """Prompt the user to enter a new account id."""
        while True:
            newid = raw_input("Enter an account number for %s: " %
                              repr(account_name))
            try:
                newid = int(newid)
            except ValueError:
                for row in flowrate.db.execute(
                    "SELECT id, name FROM accounts ORDER BY id").fetchall():
                    print '%s: "%s"' % (row.id, row.name)
                continue

            print "asset | liability | income | expense"
            newtype = raw_input("Enter an account type for %s=%s: " %
                                (newid, repr(account_name)))

            try:
                flowrate.db.execute(
                    "INSERT INTO accounts (id, name, type) VALUES (%s, %s, %s)",
                    (newid, account_name, newtype))
            except sqlalchemy.exc.IntegrityError:
                a = flowrate.db.execute(
                    "SELECT name FROM accounts WHERE id = %s",
                    (newid,)).fetchone()
                print "That account number is already in use by %s" % repr(a.name)
                continue
            else:
                return {"id": newid, "name": account_name, "type": newtype}

    def add_tx(self, credit, debit, amount, postdate, description):
        MintImporter.add_tx(self, credit, debit, amount, postdate, description)
        print credit['id'], debit['id'], amount, postdate, description


if __name__ == '__main__':
    import sys
    try:
        csvfilename, configfilename = sys.argv[1:]
    except ValueError:
        print "csvutil requires a CSV file and config file arguments."
        sys.exit(1)

    config = simplejson.loads(open(configfilename, 'rb').read())
    flowrate.set_db(
        'postgresql://%(user)s:%(password)s@%(host)s:%(port)s/%(database)s' %
        config['db_info'], echo=False, max_overflow=10, strategy="threadlocal")

    csvfile = open(csvfilename, 'rb')
    CommandLineMintImporter().load_mint_csv(csvfile)