DrCr/drcr/statements/importers/ofx2.py

#   DrCr: Web-based double-entry bookkeeping framework
#   Copyright (C) 2022–2024  Lee Yingtong Li (RunasSudo)
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

from ..models import StatementLine, reporting_commodity

from datetime import datetime
from io import StringIO
import xml.etree.ElementTree as ET

def import_ofx2(file):
	raw_ofx = file.read().decode('utf-8')
	
	# Convert OFX header to XML and parse
	xml_header = '<?xml version="1.0" encoding="UTF-8" standalone="no"?>'
	raw_payload = raw_ofx[raw_ofx.index('?>')+2:]
	xml_input = StringIO(xml_header + raw_payload.replace('&', '&amp;'))
	try:
		tree = ET.parse(xml_input)
	except Exception as ex:
		raise ex
	root = tree.getroot()
	
	# Read transactions
	lines = []  # Do first pass to catch "extra description lines"
	for transaction in root.find('BANKMSGSRSV1').find('STMTTRNRS').find('STMTRS').find('BANKTRANLIST').findall('STMTTRN'):
		date = transaction.find('DTPOSTED').text
		date = date[0:4] + '-' + date[4:6] + '-' + date[6:8]
		description = transaction.find('NAME').text
		amount = transaction.find('TRNAMT').text
		
		if amount == '0':
			lines[-1][3].append(description)
			continue
		
		lines.append([date, description, amount, []])
	
	imported_statement_lines = []
	
	# Import
	for date, description, amount, notes in lines:
		imported_statement_lines.append(StatementLine(
			dt=datetime.strptime(date, '%Y-%m-%d'),
			description=description,
			quantity=round(float(amount)*100),
			commodity=reporting_commodity()
		))
	
	return imported_statement_lines