Skip to content

Commit

Permalink
Fixes #451 and #457
Browse files Browse the repository at this point in the history
  • Loading branch information
Rahul Amaram authored and Rahul Amaram committed Jan 22, 2024
1 parent 339b854 commit 4ab38be
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 78 deletions.
130 changes: 65 additions & 65 deletions cgt_calc/parsers/schwab.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,11 @@
"""Charles Schwab parser."""
from __future__ import annotations

from collections import defaultdict
from collections import OrderedDict, defaultdict
import csv
from dataclasses import dataclass
import datetime
from decimal import Decimal
import itertools
from pathlib import Path

from cgt_calc.const import TICKER_RENAMES
Expand All @@ -18,7 +17,6 @@
)
from cgt_calc.model import ActionType, BrokerTransaction


@dataclass
class AwardPrices:
"""Class to store initial stock prices."""
Expand Down Expand Up @@ -106,32 +104,48 @@ class SchwabTransaction(BrokerTransaction):

def __init__(
self,
row: list[str],
row_dict: OrderedDict[str, str],
file: str,
):
"""Create transaction from CSV row."""
if len(row) < 8 or len(row) > 9:
if len(row_dict) < 8 or len(row_dict) > 9:
# Old transactions had empty 9th column.
raise UnexpectedColumnCountError(row, 8, file)
if len(row) == 9 and row[8] != "":
raise UnexpectedColumnCountError(list(row_dict.values()), 8, file)
if len(row_dict) == 9 and list(row_dict.values())[8] != "":
raise ParsingError(file, "Column 9 should be empty")
as_of_str = " as of "
if as_of_str in row[0]:
index = row[0].find(as_of_str) + len(as_of_str)
date_str = row[0][index:]
if as_of_str in row_dict["Date"]:
index = row_dict["Date"].find(as_of_str) + len(as_of_str)
date_str = row_dict["Date"][index:]
else:
date_str = row[0]
date_str = row_dict["Date"]
date = datetime.datetime.strptime(date_str, "%m/%d/%Y").date()
self.raw_action = row[1]
self.raw_action = row_dict["Action"]
action = action_from_str(self.raw_action)
symbol = row[2] if row[2] != "" else None
symbol = row_dict["Symbol"] if row_dict["Symbol"] != "" else None
if symbol is not None:
symbol = TICKER_RENAMES.get(symbol, symbol)
description = row[3]
quantity = Decimal(row[4].replace(",", "")) if row[4] != "" else None
price = Decimal(row[5].replace("$", "")) if row[5] != "" else None
fees = Decimal(row[6].replace("$", "")) if row[6] != "" else Decimal(0)
amount = Decimal(row[7].replace("$", "")) if row[7] != "" else None
description = row_dict["Description"]
price = (
Decimal(row_dict["Price"].replace("$", ""))
if row_dict["Price"] != ""
else None
)
quantity = (
Decimal(row_dict["Quantity"].replace(",", ""))
if row_dict["Quantity"] != ""
else None
)
fees = (
Decimal(row_dict["Fees & Comm"].replace("$", ""))
if row_dict["Fees & Comm"] != ""
else Decimal(0)
)
amount = (
Decimal(row_dict["Amount"].replace("$", ""))
if row_dict["Amount"] != ""
else None
)

currency = "USD"
broker = "Charles Schwab"
Expand All @@ -150,10 +164,10 @@ def __init__(

@staticmethod
def create(
row: list[str], file: str, awards_prices: AwardPrices
row_dict: OrderedDict[str, str], file: str, awards_prices: AwardPrices
) -> SchwabTransaction:
"""Create and post process a SchwabTransaction."""
transaction = SchwabTransaction(row, file)
transaction = SchwabTransaction(row_dict, file)
if (
transaction.price is None
and transaction.action == ActionType.STOCK_ACTIVITY
Expand All @@ -179,23 +193,7 @@ def read_schwab_transactions(
try:
with Path(transactions_file).open(encoding="utf-8") as csv_file:
lines = list(csv.reader(csv_file))

headers = [
"Date",
"Action",
"Symbol",
"Description",
"Quantity",
"Price",
"Fees & Comm",
"Amount",
]
if not lines[0] == headers:
raise ParsingError(
transactions_file,
"First line of Schwab transactions file must be something like "
"'Transactions for account ...'",
)
headers = lines[0]

if len(lines[1]) < 8 or len(lines[1]) > 9:
raise ParsingError(
Expand All @@ -204,16 +202,12 @@ def read_schwab_transactions(
" with 8 columns",
)

if "Total" not in lines[-1][0]:
raise ParsingError(
transactions_file,
"Last line of Schwab transactions file must be total",
)

# Remove headers and footer
lines = lines[1:-1]
# Remove header
lines = lines[1:]
transactions = [
SchwabTransaction.create(row, transactions_file, awards_prices)
SchwabTransaction.create(
OrderedDict(zip(headers, row)), transactions_file, awards_prices
)
for row in lines
]
transactions.reverse()
Expand All @@ -229,15 +223,19 @@ def _read_schwab_awards(
"""Read initial stock prices from CSV file."""
initial_prices: dict[datetime.date, dict[str, Decimal]] = defaultdict(dict)

headers = []

lines = []
if schwab_award_transactions_file is not None:
try:
with Path(schwab_award_transactions_file).open(
encoding="utf-8"
) as csv_file:
lines = list(csv.reader(csv_file))
headers = lines[0]

# Remove headers
lines = lines[2:]
lines = lines[1:]
except FileNotFoundError:
print(
"WARNING: Couldn't locate Schwab award "
Expand All @@ -246,36 +244,38 @@ def _read_schwab_awards(
else:
print("WARNING: No schwab award file provided")

modulo = len(lines) % 3
modulo = len(lines) % 2
if modulo != 0:
raise UnexpectedRowCountError(
len(lines) - modulo + 3, schwab_award_transactions_file or ""
len(lines) - modulo + 2, schwab_award_transactions_file or ""
)

for row in zip(lines[::3], lines[1::3], lines[2::3]):
if len(row) != 3:
raise UnexpectedColumnCountError(
list(itertools.chain(*row)), 3, schwab_award_transactions_file or ""
)

lapse_main, _, lapse_data = row
for upper_row, lower_row in zip(lines[::2], lines[1::2]):
# in this format each row is split into two rows,
# so we combine them safely below
row = []
for upper_col, lower_col in zip(upper_row, lower_row):
assert upper_col == "" or lower_col == ""
row.append(upper_col + lower_col)

if len(lapse_main) != 8:
raise UnexpectedColumnCountError(
lapse_main, 8, schwab_award_transactions_file or ""
)
if len(lapse_data) < 8 or len(lapse_data) > 9:
if len(row) != len(headers):
raise UnexpectedColumnCountError(
lapse_data, 8, schwab_award_transactions_file or ""
row, len(headers), schwab_award_transactions_file or ""
)

date_str = lapse_main[0]
row_dict = OrderedDict(zip(headers, row))
date_str = row_dict["Date"]

try:
date = datetime.datetime.strptime(date_str, "%Y/%m/%d").date()
except ValueError:
date = datetime.datetime.strptime(date_str, "%m/%d/%Y").date()
symbol = lapse_main[2] if lapse_main[2] != "" else None
price = Decimal(lapse_data[3].replace("$", "")) if lapse_data[3] != "" else None
symbol = row_dict["Symbol"] if row_dict["Symbol"] != "" else None
price = (
Decimal(row_dict["FairMarketValuePrice"].replace("$", ""))
if row_dict["FairMarketValuePrice"] != ""
else None
)
if symbol is not None and price is not None:
symbol = TICKER_RENAMES.get(symbol, symbol)
initial_prices[date][symbol] = price
Expand Down
25 changes: 12 additions & 13 deletions tests/test_data/schwab_transactions.csv
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
"Date","Action","Symbol","Description","Quantity","Price","Fees & Comm","Amount"
"04/02/2021","Buy","FOO","FOO INC","30.5","$30.2","$4","-$925.1"
"03/06/2021","Sell","FOO","FOO INC","90","$33","$5","$2965"
"03/06/2021","Buy","FOO","FOO INC","90","$32","$5","-$2885"
"03/03/2021","Sell","FOO","FOO INC","104","$31","$5","$3219"
"03/03/2021","Buy","FOO","FOO INC","4","$30","$0","-$120"
"03/03/2021","Buy","FOO","FOO INC","100","$29","$5","-$2905"
"03/03/2021","Sell","FOO","FOO INC","50","$28","$5","$1395"
"03/03/2021","Buy","FOO","FOO INC","50","$27","$5","-$1355"
"03/03/2021","Sell","FOO","FOO INC","100","$26","$5","$2595"
"03/02/2021","Buy","FOO","FOO INC","100","$25","$6","-$2506"
"03/01/2016","MoneyLink Transfer","","Tfr BANK","","","","$10000.00"
Transactions Total,"","","","","","",
Date,Action,Symbol,Description,Price,Quantity,Fees & Comm,Amount
04/02/2021,Buy,FOO,FOO INC,$30.2,30.5,$4,-$925.1
03/06/2021,Sell,FOO,FOO INC,$33,90,$5,$2965
03/06/2021,Buy,FOO,FOO INC,$32,90,$5,-$2885
03/03/2021,Sell,FOO,FOO INC,$31,104,$5,$3219
03/03/2021,Buy,FOO,FOO INC,$30,4,$0,-$120
03/03/2021,Buy,FOO,FOO INC,$29,100,$5,-$2905
03/03/2021,Sell,FOO,FOO INC,$28,50,$5,$1395
03/03/2021,Buy,FOO,FOO INC,$27,50,$5,-$1355
03/03/2021,Sell,FOO,FOO INC,$26,100,$5,$2595
03/02/2021,Buy,FOO,FOO INC,$25,100,$6,-$2506
03/01/2016,MoneyLink Transfer,,Tfr BANK,,,,$10000.00

0 comments on commit 4ab38be

Please sign in to comment.