Fetch prices from IEX API for Beancount

fetch-prices.py

#!/usr/bin/env python

'''A quick-and-dirty script to fetch prices from IEX API.'''

import sys
import argparse

import requests
from dateutil.parser import parse as parse_datetime
from beancount.core import data
from beancount.core.number import D
from beancount.core.amount import Amount
from beancount.parser import printer
from beancount.loader import load_file
from beancount.ops import holdings


# https://www.iexcloud.io/cloud-login#/register
# Sign up for a free account and put your public token here
IEX_TOKEN = ''


def get_prices_from_iex(symbol, date_range='3m'):

    # Message usage estimate:
    # 1 symbol * 3 months chartCloseOnly data is about 128 messages
    # Free account monthly quota: 500,000

    print('Fetching prices for {} ...'.format(symbol), file=sys.stderr)

    url = f'https://cloud.iexapis.com/stable/stock/{symbol}/chart/{date_range}'
    params = {
        'chartCloseOnly': True,
        'token': IEX_TOKEN,
    }

    json_data = requests.get(url, params=params).json()
    prices = [
        (symbol, parse_datetime(item['date']).date(), item['close'])
        for item in json_data
    ]
    return prices


def convert_prices(prices):
    '''Convert prices to Beancount "Price" entries'''
    meta = data.new_metadata('<price>', 0)
    entries = [
        data.Price(
            meta=meta,
            date=price[1],
            currency=price[0],
            amount=Amount(D('{:.2f}'.format(price[2])), 'USD')
        )
        for price in prices
    ]
    return entries


def extract_symbols(filename):
    entries, _, _ = load_file(filename)
    symbols = set(
        h.currency for h in holdings.get_final_holdings(entries)
        if h.account.startswith('Assets:') and 'Positions' in h.account
    )
    return symbols


def fetch_all(fetch_func, symbols):
    symbols = symbols or []
    entries = []
    for symbol in symbols:
        prices = fetch_func(symbol)
        _entries = convert_prices(prices)
        entries.extend(_entries)
    entries.sort()
    return entries


def main():
    argparser = argparse.ArgumentParser()
    symbol_source = argparser.add_mutually_exclusive_group(required=True)
    symbol_source.add_argument('-f', '--from-file', help='read symbols from Beancount file')
    symbol_source.add_argument('-s', '--symbols', metavar='SYMBOL', nargs='+', help='read symbols from command line')
    argparser.add_argument('-l', '--list', action='store_true', help='only list symbols to fetch prices for')
    args = argparser.parse_args()

    if args.from_file:
        symbols = extract_symbols(args.from_file)
    elif args.symbols:
        symbols = args.symbols

    if not args.list:
        entries = fetch_all(get_prices_from_iex, symbols)
        printer.print_entries(entries)
    else:
        print(symbols)


if __name__ == '__main__':
    main()