ynrc/src/ynrc/main.py

# -*- coding: utf-8 -*-
import argparse
import logging
import sys
import datetime
import json

import hammock
import requests

from ynrc import __version__

__author__ = "Nikola Kotur"
__copyright__ = "Nikola Kotur"
__license__ = "mit"

_logger = logging.getLogger(__name__)


def parse_args(args):
    parser = argparse.ArgumentParser(description="You Need Rocket.Chat: send weekly YNAB reports to your Rocket.Chat instance")
    parser.add_argument(
        "--version",
        action="version",
        version="ynrc {ver}".format(ver=__version__),
    )
    parser.add_argument("--ynab-api-key", help="YNAB API key", type=str, required=True)
    parser.add_argument("--ynab-budget-name", help="YNAB budget name", type=str, required=True)
    parser.add_argument("--rocketchat-webhook", help="Rocket.Chat webhook URL", type=str, required=True)
    parser.add_argument("--rocketchat-avatar-url", help="Rocket.Chat user avatar URL", type=str, default="")
    parser.add_argument(
        "-v",
        "--verbose",
        dest="loglevel",
        help="set loglevel to INFO",
        action="store_const",
        const=logging.INFO,
    )
    parser.add_argument(
        "-vv",
        "--very-verbose",
        dest="loglevel",
        help="set loglevel to DEBUG",
        action="store_const",
        const=logging.DEBUG,
    )
    return parser.parse_args(args)


def setup_logging(loglevel):
    logformat = "[%(asctime)s] %(levelname)s:%(name)s:%(message)s"
    logging.basicConfig(
        level=loglevel, stream=sys.stdout, format=logformat, datefmt="%Y-%m-%d %H:%M:%S"
    )


def _get_group_name_by_id(groups, group_id):
    for group in groups:
        if group_id == group['id']:
            return group['name']
    return ''


def fetch_from_ynab(api_key, budget_name):
    _logger.debug("Fetching YNAB budget %s...", budget_name)
    ynab = hammock.Hammock(
        "https://api.youneedabudget.com/v1",
        headers={"Authorization": "Bearer %s" % api_key},
    )

    budgets = ynab.budgets.GET().json()['data']['budgets']
    for budget in budgets:
        if budget['name'] == budget_name:
            budget_id = budget['id']
            budget_full = ynab.budgets(budget_id).GET().json()['data']['budget']
            category_groups = budget_full['category_groups']
            break
    else:
        raise Exception('Budget with name "%s" not found' % budget_name)

    _logger.debug('Using %s (%s)', budget_name, budget_id)

    today = datetime.date.today()
    week_ago = today - datetime.timedelta(days=7)
    since_date = week_ago.strftime('%Y-%m-%d')

    _logger.debug('Fetching all transactions since %s...', since_date)

    transactions = ynab.budgets(budget_id).transactions.GET('?since_date=%s' % since_date).json()['data']['transactions']

    # Calculate spending
    res = {}
    for tx in transactions:
        if tx['subtransactions']:
            tx = tx['subtransactions']
        else:
            tx = [tx]
        for t in tx:
            if t['amount'] > 0:
                continue
            tid = t['category_id']
            if tid not in res:
                res[tid] = {
                    'amount': 0.0,
                }
            res[tid]['name'] = t.get('category_name', None)
            res[tid]['amount'] += t['amount']

    # Get category names
    for tid, t in res.items():
        if tid is None:
            res[tid]['name'] = 'Transfer'
            continue
        cat = ynab.budgets(budget_id).categories(tid).GET().json()['data']['category']
        parent_cat = _get_group_name_by_id(category_groups, cat['category_group_id'])
        res[tid]['name'] = parent_cat + ': ' + cat['name']

    # Sort top 10 results
    out = []
    for r in (sorted(res.values(), key=lambda k: k['amount'])):
        amount = u'%4.2f €' % abs(r['amount'] / 1000)
        out.append((u'%s' % r['name'], amount, ))
        if len(out) >= 10:
            break
    return out


def post_to_rocketchat(hook, avatar, ynab):
    if not ynab:
        return

    today = datetime.date.today()
    week_ago = today - datetime.timedelta(days=7)

    post_data = {
        "alias": "YNAB",
        "text": 'Report for week of %s (from %s)' % (today.strftime('%Y-%W'), week_ago.strftime('%a %Y-%m-%d')),
        "attachments": [
        ],
    }
    if avatar:
        post_data['avatar'] = avatar

    fields = []
    for spending in ynab:
        spending_name, spending_amount = spending
        fields.append({
            "title": spending_name,
            "value": spending_amount,
        })
    post_data['attachments'] = [{
        'title': 'Top 10 categories spent',
        'fields': fields,
    }]

    result = requests.post(
        hook,
        headers={'Content-Type': 'application/json'},
        data=json.dumps(post_data),
    )
    if result.json().get('success', False):
        _logger.debug("Posted to Rocket.Chat successfully")
    else:
        _logger.error("Failed posting to Rocket.Chat: %s". result.text)


def main_call(args):
    args = parse_args(args)
    setup_logging(args.loglevel)
    _logger.debug("Starting ynrc...")

    ynab = fetch_from_ynab(args.ynab_api_key, args.ynab_budget_name)
    post_to_rocketchat(args.rocketchat_webhook, args.rocketchat_avatar_url, ynab)

    _logger.debug("Finished running ynrc")


def run():
    """Entry point for console_scripts"""
    main_call(sys.argv[1:])


if __name__ == "__main__":
    run()