ledger_amounts.py

#!/usr/bin/env python

import subprocess
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from typing import Optional

import common

LEDGER_COMMODITY_CMD = (
    f"{common.LEDGER_BIN} -f {common.LEDGER_DAT} "
    '--balance-format "%(S(display_total))" -n -c bal'
)
LEDGER_LIMIT_ETFS = (
    f"""--limit 'commodity!~/{common.CURRENCIES_REGEX}/ and commodity=~/^[A-Z]+/'"""
)


def get_commodity_amounts(ledger_args: str) -> dict[str, float]:
    process = subprocess.run(
        f"{LEDGER_COMMODITY_CMD} {ledger_args}",
        shell=True,
        check=True,
        text=True,
        capture_output=True,
    )

    lines = process.stdout.splitlines()
    df_data = {}
    for line in lines:
        shares, ticker = line.split(maxsplit=1)
        ticker = ticker.strip('"')
        df_data[ticker] = float(shares)
    return df_data


def get_etfs_amounts(account: Optional[str] = None) -> dict[str, float]:
    with ThreadPoolExecutor() as exc:
        results = []
        if account:
            results.append(
                exc.submit(
                    get_commodity_amounts,
                    LEDGER_LIMIT_ETFS
                    + f' --limit "account=~/^Assets:Investments:{account}/"',
                )
            )
        else:
            results.append(
                exc.submit(
                    get_commodity_amounts,
                    LEDGER_LIMIT_ETFS
                    + ' --limit "account=~/^Assets:Investments:(Charles Schwab .*Brokerage|Interactive Brokers)/"',
                )
            )
            results.append(
                exc.submit(
                    get_commodity_amounts,
                    LEDGER_LIMIT_ETFS
                    + ' --limit "account=~/^Assets:Investments:Retirement:Charles Schwab IRA/"',
                )
            )
    combined = defaultdict(float)
    for r in results:
        for k, v in r.result().items():
            combined[k] += v
    return combined