example_bank_etl.py

#354
Raw
Author
winny
Created
April 18, 2021, 6:39 a.m.
Expires
Never
Size
3.4 KB
Hits
289
Syntax
Python
Private
✗ No
#!/usr/bin/env python3
from abc import abstractmethod, ABC
from datetime import datetime
import csv
import argparse
import sys


def date_type(s):
    """Parse a date, used for the add_argument()'s type keyword argument."""
    if s in ('today', 'now'):
        return datetime.today()
    try:
        return datetime.strptime(s, '%Y-%m-%d')
    except ValueError:
        raise argparse.ArgumentTypeError(
            f"`{s}' invalid format. Try now, today, or a ISO date (%Y-%m-%d)."
        )


def format_type(s):
    """Parse the csv format, used for add_argument()'s type keyword argument.
    """
    if s not in ('credit', 'checking'):
        raise argparse.ArgumentTypeError(
            f"`{s}' invalid.  Try `checking' or `credit'."
        )
    return s


def build_parser():
    """Build a parser and return it."""
    parser = argparse.ArgumentParser(
        description='Transform a Bank of XXXXXX csv to something ledger'
                    ' can use.  Try importing it into ledge as follows:'
                    ' ledger -f your.ledger convert --account'
                    ' Assets:Bank:Checking massaged.csv'
    )

    parser.add_argument(
        '--begin',
        type=date_type,
        default=datetime.min,
        help='now, today, or a ISO date (%%Y-%%m-%%d). '
             'Defaults to earliest date possible.'
    )

    parser.add_argument(
        '--format',
        type=format_type,
        required=True,
        help='checking or credit',
    )

    return parser


class Translation(ABC):
    @staticmethod
    def flip(s):
        """Flip the sign of a string currency value."""
        if s.startswith('-'):
            return s[1:]
        else:
            return f'-{s}'

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

    @abstractmethod
    def date(self):
        pass

    @abstractmethod
    def amount(self):
        pass

    @abstractmethod
    def description(self):
        pass

    @abstractmethod
    def code(self):
        pass


class CheckingTranslation(Translation):
    def date(self):
        return datetime.strptime(self.row['Date'], '%m/%d/%Y')

    def amount(self):
        if self.row['Credit']:
            return f'-{self.row["Credit"]}'
        return self.row['Debit']

    def description(self):
        return self.row['Description']

    def code(self):
        return self.row['ChkRef']


class CreditTranslation(Translation):
    def date(self):
        return datetime.strptime(self.row['PostDate'], '%m/%d/%y')

    def amount(self):
        return Translation.flip(self.row['Amount'])

    def description(self):
        return self.row['Description']

    def code(self):
        return self.row['Reference'].strip("'")


TRANSLATION = {
    'checking': CheckingTranslation,
    'credit': CreditTranslation,
}


def main():
    """ETL tool entry point."""
    args = build_parser().parse_args()

    Translation = TRANSLATION[args.format]

    reader = csv.DictReader(sys.stdin)
    writer = csv.DictWriter(sys.stdout,
                            ('date', 'amount', 'code', 'description'))
    writer.writeheader()
    for r in reader:
        t = Translation(r)
        date = t.date()
        if date < args.begin:
            continue
        writer.writerow({
            'date': date.strftime('%Y-%m-%d'),
            'code': t.code(),
            'description': t.description(),
            'amount': t.amount(),
        })


if __name__ == '__main__':
    main()