orgtool / orgtool / ext / finances /

# -*- coding: utf-8 -*-
import datetime
from decimal import Decimal as D
from argh import *
from tool import app

from .schema import Contract, Payment

DATA_SOURCE = __name__

# FIXME refactor this function (it's almost everywhere!)
def args_to_unicode(namespace, names):
    for name in names:
        value = getattr(namespace, name, None)
        if value is not None:
            setattr(namespace, name, value.decode('utf8'))

def get_contracts(query):
    db = app.get_feature('document_storage').default_db
    contracts = Contract.objects(db)
    if query:
        return contracts.where(summary__matches_caseless=query)
        return contracts

def get_payments(query=None, contract_query=None):
    db = app.get_feature('document_storage').default_db
    contracts = get_contracts(contract_query) if contract_query else None
    payments = Payment.objects(db)
    if query:
        payments = payments.where(summary__matches_caseless=query)
    if contract_query:
        if not contracts:
            return []
        payments = payments.where(plan__in=[ for c in contracts])
    return payments

@arg('-q', '--query')
def list_contracts(args):
    "Displays a list of matching contracts."
    args_to_unicode(args, ['query'])
    contracts = get_contracts(args.query)
    for c in contracts:
        yield unicode(c)

@arg('-q', '--query')
@arg('-c', '--contract')
def list_payments(args):
    "Displays a list of matching payments."
    args_to_unicode(args, ['query', 'contract'])
    payments = get_payments(args.query, args.contract)
    for p in payments:
        yield unicode(p)

@arg('amount', help='e.g. "-9.99" (negative = expences, positive = income)')
@arg('currency', help='e.g. "USD"')
@arg('-d', '--date', help='e.g. 2010-12-27 (leave empty for today)')
@arg('-t', '--time', help='e.g. 14:52 (leave empty for 00:00)')
@arg('-c', '--contract', help='contract (query on summary)')
@arg('--balance', help='balance after the payment')
@arg('-s', '--summary', help='usually not needed when a contract is specified')
def add_payment(args):
    "Logs a payment with given properties."
    args_to_unicode(args, ['currency', 'summary', 'contract'])
    contract = None
    if args.contract:
        contracts = get_contracts(args.contract)
        if not contracts:
            raise COmmandError('no contracts matching query')
        if 1 < len(contracts):
            yield 'Mare than one contract matched query:'
            for c in contracts:
                yield unicode(c)
            raise CommandError('Please use a more precise query')

        contract = contracts[0]

    date_time = datetime.datetime.utcnow()
        year, month, day = [int(x) for x in'-')]
        h, m = [int(x) for x in args.time.split(':')] if args.time else (0, 0)
        date_time = datetime.datetime(year, month, day, h, m)

    amount = D(args.amount)
    balance = D(args.balance) if args.balance else None

    payment = Payment(
        date_time = date_time,
        amount = amount,
        currency = args.currency,
        plan = contract,
        summary = args.summary,
        balance = balance,
        source = DATA_SOURCE+'.add_payment',

    if not args.dry_run:
        db = app.get_feature('document_storage').default_db

    yield u'Added payment:'
    yield ''
    for k, v in payment.iteritems():
        if v is not None:
            yield u'  {0}: {1}'.format(k, v)
    yield ''
    yield u'assigned primary key {0}'.format(

command_list = list_contracts, list_payments, add_payment