Determine the cost basis of a robo-traded stock from Wealthfront exported as Quicken QFX

cost_basis.py

#!/usr/bin/env python3
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from collections import namedtuple
from dataclasses import dataclass
import csv
from typing import Sequence
import operator

__author__ = 'David Blume'
__copyright__ = "Copyright 2024, David Blume"
__license__ = "Apache 2.0"
__version__ = "1.0"


@dataclass
class Tranche:
    date: str
    units: float
    unit_price: float


transaction_fields = ('type',
                      'trade_date',
                      'units',
                      'unit_price',
                      'security_name',
                      'security_ticker')


Transaction = namedtuple('Transaction', transaction_fields)


def lifo(tranches: Sequence[Tranche]) -> int:
    """LIFO: Last In First Out"""
    return -1


def fifo(tranches: Sequence[Tranche]) -> int:
    """FIFO: First In First Out"""
    return 0


def hifo(tranches: Sequence[Tranche]) -> int:
    """HIFO: Highest (unit price) In First Out"""
    max_index, max_value = max(enumerate(tranches), key=lambda x: x[1].unit_price)
    return max_index


def main(strategy: str, ticker: str, filename: str) -> None:
    transactions = list()
    with open(filename) as f:
        reader = csv.reader(f)
        headers = next(reader)
        # Supports the columns from convert_qfx_to_csv.py
        # type,tradeDate,settleDate,memo,income_type,units,unit_price,comission,fees,total,tferaction,securityid,securityname,securityticker,securitymemo
        assert headers[0] == 'type' and headers[1] == 'tradeDate' and \
               headers[5] == 'units' and headers[6] == 'unit_price' and \
               headers[13] == 'securityticker'
        for r in reader:
            if r[13] == ticker:
                for i in (5, 6):
                    if r[i] == '':
                        r[i] = '0'
                transaction = Transaction(r[0], r[1], *map(float, r[5:7]), r[12], r[13])
                transactions.append(transaction)

    # Sort chronologically
    transactions.sort(key=lambda x: x.trade_date)

    # Select the strategy
    if strategy == 'lifo':
        select_tranche = lifo
    elif strategy == 'fifo':
        select_tranche = fifo
    elif strategy == 'hifo':
        select_tranche = hifo

    tranches = []
    for t in transactions:
        if t.type == 'buystock':
            tranches.append(Tranche(t.trade_date, t.units, t.unit_price))
        elif t.type == 'sellstock':
            units = -t.units
            while units:
                tranche_id = select_tranche(tranches)
                tranch = tranches[tranche_id]
                if units >= tranch.units:
                    units -= tranch.units
                    tranches.pop(tranche_id)
                else:
                    tranch.units -= units
                    units = 0

    units_sum = sum([t.units for t in tranches])
    price_sum = sum([t.unit_price * t.units for t in tranches])
    print(f"{strategy.upper()} strategy left {units_sum:.0f} shares worth ${price_sum:.2f} in {len(tranches)} tranches.")
    for i in tranches:
        print(f'{i.units:.0f} for ${i.unit_price * i.units:.2f} on {i.date.split(" ")[0]}')


if __name__ == '__main__':
    parser = ArgumentParser(formatter_class=RawDescriptionHelpFormatter,
                            description='''
Determine cost basis for a stock, given a CSV of transactions.
The CSV file should have been generated from a Quicken QFX file via
https://gist.github.com/dblume/946f6d88701371eec07d9e8c0d8432f6

Available selling strategies:
* LIFO: Last In First Out
* FIFO: First In First Out
* HIFO: Highest In First Out
''')
    parser.add_argument('--strategy', choices=['lifo', 'fifo', 'hifo'], default='lifo', help='The sell algorithm.')
    parser.add_argument('ticker', help='stock ticker symbol')
    parser.add_argument('filename', help='CSV file of transactions.')
    args = parser.parse_args()
    main(args.strategy, args.ticker.upper(), args.filename)