From 8ef62592bfb3c34c1098c1801e807af92efebbe9 Mon Sep 17 00:00:00 2001 From: Jonas Trappenberg Date: Sat, 4 Mar 2023 12:24:23 +0100 Subject: [PATCH] Key `Account.transactions` by account type instead of str --- quiffen/core/account.py | 7 +++---- tests/test_account.py | 8 +++---- tests/test_qif.py | 46 +++++++++++++++++++++++++++++------------ 3 files changed, 40 insertions(+), 21 deletions(-) diff --git a/quiffen/core/account.py b/quiffen/core/account.py index 8f45645..ad4ce49 100644 --- a/quiffen/core/account.py +++ b/quiffen/core/account.py @@ -40,7 +40,7 @@ class Account(BaseModel): >>> acc Account(name='Example name', desc='Some description') >>> tr = quiffen.Transaction(date=datetime.now(), amount=150.0) - >>> acc.set_last_header = quiffen.AccountType.BANK + >>> acc.set_header(quiffen.AccountType.BANK) >>> acc.add_transaction(tr) >>> acc.transactions {'Bank': [Transaction(date=datetime.datetime(2021, 7, 2, 18, 31, 47, 817025), amount=150.0)]} @@ -60,7 +60,7 @@ class Account(BaseModel): credit_limit: Decimal = None balance: Decimal = None date_at_balance: datetime = None - transactions: Dict[str, TransactionList] = {} + transactions: Dict[AccountType, TransactionList] = {} _last_header: AccountType = None __CUSTOM_FIELDS: List[Field] = [] @@ -123,7 +123,6 @@ def add_transaction( RuntimeError If there is no header provided, and no ``last_header`` set. """ - if not header and not self._last_header: raise RuntimeError('No header provided, and no last header set.') @@ -199,7 +198,7 @@ def to_qif( ) for header, header_transactions in self.transactions.items(): - qif += f'^\n!Type:{header}\n' + qif += f'^\n!Type:{header.value}\n' qif += '^\n'.join( transaction.to_qif( date_format=date_format, diff --git a/tests/test_account.py b/tests/test_account.py index eaa2ff1..c60a1a8 100644 --- a/tests/test_account.py +++ b/tests/test_account.py @@ -3,7 +3,7 @@ import pytest -from quiffen import Account, Transaction +from quiffen import Account, AccountType, Transaction def test_create_account(): @@ -83,7 +83,7 @@ def test_str_method(): account.add_transaction( Transaction(date=datetime(2022, 2, 1), amount=0), - header='Bank', + header=AccountType.BANK, ) assert str(account) == ( 'Account:\n\tName: Test Account\n\tDesc: Test Description\n\t' @@ -92,7 +92,7 @@ def test_str_method(): account.add_transaction( Transaction(date=datetime(2022, 2, 1), amount=0), - header='CCard', + header=AccountType.CREDIT_CARD, ) assert str(account) == ( 'Account:\n\tName: Test Account\n\tDesc: Test Description\n\t' @@ -105,7 +105,7 @@ def test_set_header(): """Test setting the header for an account""" account = Account(name='Test Account') assert account._last_header is None - account.set_header('Bank') + account.set_header(AccountType.BANK) assert account._last_header == 'Bank' # pylint: enable=protected-access diff --git a/tests/test_qif.py b/tests/test_qif.py index cee4266..81f9f29 100644 --- a/tests/test_qif.py +++ b/tests/test_qif.py @@ -547,7 +547,7 @@ def test_get_data_dicts_transactions(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) qif.add_account(account) data_dicts = qif._get_data_dicts(data_type=QifDataType.TRANSACTIONS) @@ -570,7 +570,7 @@ def test_get_data_dicts_transactions_with_date_format_and_ignore(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) qif.add_account(account) data_dicts = qif._get_data_dicts( @@ -613,7 +613,7 @@ def test_get_data_dicts_accounts(): """Test the get_data_dicts method with accounts""" qif = Qif() account = Account(name='Test Account') - account.set_header('Bank') + account.set_header(AccountType.BANK) qif.add_account(account) data_dicts = qif._get_data_dicts(data_type=QifDataType.ACCOUNTS) @@ -633,7 +633,7 @@ def test_get_data_dicts_investments(): security='Test Security', price=Decimal('10'), ) - account.set_header('Invst') + account.set_header(AccountType.INVST) account.add_transaction(investment) qif.add_account(account) data_dicts = qif._get_data_dicts(data_type=QifDataType.INVESTMENTS) @@ -667,7 +667,7 @@ def test_to_csv_transactions(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) qif.add_account(account) @@ -734,7 +734,7 @@ def test_to_csv_accounts(): """Test the to_csv method with accounts""" qif = Qif() account = Account(name='Test Account', account_type='Bank') - account.set_header('Bank') + account.set_header(AccountType.BANK) qif.add_account(account) csv_file = Path(__file__).parent / 'test_files' / 'test_output.csv' @@ -762,7 +762,7 @@ def test_to_csv_investments(): security='Test Security', price=Decimal('10'), ) - account.set_header('Invst') + account.set_header(AccountType.INVST) account.add_transaction(investment) qif.add_account(account) @@ -814,7 +814,7 @@ def test_to_csv_transactions_with_date_format_and_ignore_list(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) qif.add_account(account) @@ -852,7 +852,7 @@ def test_to_csv_transactions_multiple(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) transaction2 = transaction.copy() @@ -895,7 +895,7 @@ def test_to_dataframe_transactions(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) qif.add_account(account) @@ -973,7 +973,7 @@ def test_to_dataframe_investments(): security='Test Security', price=Decimal('10'), ) - account.set_header('Invst') + account.set_header(AccountType.INVST) account.add_transaction(investment) qif.add_account(account) @@ -997,7 +997,7 @@ def test_to_dataframe_transactions_with_ignore_list(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) qif.add_account(account) @@ -1025,7 +1025,7 @@ def test_to_dataframe_transactions_multiple(): memo='Test Memo', category=Category(name='Test Category'), ) - account.set_header('Bank') + account.set_header(AccountType.BANK) account.add_transaction(transaction) transaction2 = transaction.copy() @@ -1085,3 +1085,23 @@ def test_transaction_before_account_definition_2(qif_file): def test_empty_qif(): qif = Qif() assert qif.to_qif() == '' + + +def test_transaction_account_type_qif(): + qif = Qif() + account = Account(name='Test Account') + qif.add_account(account) + transaction = Transaction( + date=datetime(2019, 1, 1), + amount=Decimal('100'), + ) + account.set_header(AccountType.CASH) + account.add_transaction(transaction) + assert qif.to_qif() == '''!Account +NTest Account +^ +!Type:Cash +D2019-01-01 +T100 +^ +'''