harvest.py

import argparse
import collections
import datetime
import openpyxl
import tableformatter
import re

import tda


def extract_symbol(security):
    m = re.search('\(.+\)', security)
    assert not m.groups()
    return m.group(0)[1:-1]


def data_rows(positions_workbook):
    sec_name = positions_workbook['A2'].value
    # First row is assumed to be the header
    row = 2
    while sec_name:
        if sec_name != 'Total:':
            yield row

        row += 1
        sec_name = positions_workbook['A'+str(row)].value


def read_cell(sheet, row, col):
    return sheet[row+str(col)].value


def validate_workbook(wb):
    positions = wb['Unrealized Long']

    def expect_label(cell, name):
        value = positions[cell].value
        if value != name:
            raise ValueError(
                    'Expected cell {} to have value "{}", got "{}"'.format(
                cell, name, value))

    # Verify column labels haven't changed
    expect_label('A1', 'Security')
    expect_label('B1', 'Qty')
    expect_label('C1', 'Open date')
    expect_label('D1', 'Adj cost per share')
    expect_label('E1', 'Adj cost')
    expect_label('F1', 'Wash sale adj')
    expect_label('G1', 'Last')
    expect_label('H1', 'Mkt value')
    expect_label('I1', 'Adj gain($)')
    expect_label('J1', 'Adj gain(%)')
    expect_label('K1', 'Term')

    # Verify that we can extract tickers from security names
    for row in data_rows(positions):
        extract_symbol(positions['A'+str(row)].value)


def compute_sales_for_losses(workbook, target_term):
    assert target_term in ('long', 'short')

    positions = workbook['Unrealized Long']

    # Security, quantity, loss value
    table_row = lambda: ['', 0, 0]

    sales = collections.defaultdict(table_row)
    sales['zzzTotal'][0] = 'Total'
    sales['zzzTotal'][1] = None

    for row in data_rows(positions):
        symbol = extract_symbol(read_cell(positions, 'A', row))
        quantity = read_cell(positions, 'B', row)
        d = read_cell(positions, 'C', row)
        open_date = datetime.date(d.year, d.month, d.day)
        gain = read_cell(positions, 'I', row)

        term = read_cell(positions, 'K', row)
        if term == 'Long-term':
            term = 'long'
        elif term == 'Short-term':
            term = 'short'
        else:
            raise ValueError('Invalid term: "{}"'.format(term))

        if term == target_term and gain < 0:
            sales[symbol][0] = symbol
            sales[symbol][1] += quantity
            sales[symbol][2] += gain

            sales['zzzTotal'][2] += gain

    print(tableformatter.generate_table(
        [sales[s] for s in sorted(sales.keys())],
        (
            tableformatter.Column(col_name='Security'),
            tableformatter.Column(
                col_name='Quantity',
                formatter=
                lambda v: ('{:,.2f}'.format(v)) if v is not None else '',
                cell_halign=tableformatter.ColumnAlignment.AlignRight),
            tableformatter.Column(
                col_name='Loss',
                formatter=lambda v: '{:,.2f}'.format(v),
                cell_halign=tableformatter.ColumnAlignment.AlignRight),
        ),
        grid_style=tableformatter.FancyGrid()))


def verify_tax_lot_setting(client, account_id):
    import json
    prefs = client.get_preferences(account_id).json()

    bad_tax_lot_methods = []
    if prefs['equityTaxLotMethod'] != 'MINIMUM_TAX':
        bad_tax_lot_methods.append('Equities')
    if prefs['optionTaxLotMethod'] != 'MINIMUM_TAX':
        bad_tax_lot_methods.append('Options')
    if prefs['mutualFundTaxLotMethod'] != 'MINIMUM_TAX':
        bad_tax_lot_methods.append('Mutual Funds')

    if bad_tax_lot_methods:
        print()
        print('WARNING!!!')
        print()
        print('Your account does not enable efficient tax loss harvesting on '+
                'the following transaction types: {}.'.format(
                    ', '.join(bad_tax_lot_methods)))
        print()
        print('Continuing may cause TDAmeritrade to assign tax lots in such a '+
                'way as to defeat the purpose of this script.')
        print()
        print('You can update your tax lot ID preferences by visiting this '+
                'link and choosing "Tax efficient loss harvester:" '+
                'https://invest.ameritrade.com/grid/p/site#r=jPage/cgi-bin/'+
                'apps/u/AccountSettings')
        print()
        print('You may ignore this error by rerunning this script with '+
                '"--ignore_tax_lot_method", but be advised that you do so '+
                'AT YOUR OWN RISK.')
        print()

        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
            'Issue trades to harvest tax losses using spreadsheet of cost '+
            'basis data.')
    parser.add_argument('--token_path', type=str, required=True)
    parser.add_argument('--api_key', type=str, required=True)
    parser.add_argument('--redirect_path', type=str, required=True)
    parser.add_argument('--execute_trades', default=False)
    parser.add_argument('--gainskeeper_sheet', type=str, required=True)
    parser.add_argument('--ignore_tax_lot_method', action='store_true')
    args = parser.parse_args()

    workbook = openpyxl.load_workbook(filename=args.gainskeeper_sheet)

    client = None
    if args.token_path:
        account_id = int(input('Enter account ID to perform account actions: '))
        client = tda.auth.easy_client(args.api_key, args.redirect_path, 
                args.token_path)

    validate_workbook(workbook)
    compute_sales_for_losses(workbook, target_term='short')

    if client:
        if not args.ignore_tax_lot_method:
            verify_tax_lot_setting(client, account_id)


if __name__ == '__main__':
    main()