From 2c7c527ea937df3a04134fb080bad99531fd77a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=ADs=20Murta?= Date: Fri, 21 Apr 2023 15:32:40 +0100 Subject: [PATCH] [Refactor] Creates Tagger transformer Move the tag rules based transformer to its own class. Verified with unit tests. --- pfbudget/core/manager.py | 4 ++++ pfbudget/transform/__init__.py | 0 pfbudget/transform/categorizer.py | 31 ------------------------------- pfbudget/transform/tagger.py | 30 ++++++++++++++++++++++++++++++ tests/mocks/categories.py | 8 ++++++-- tests/test_transform.py | 18 +++++++++++++++++- 6 files changed, 57 insertions(+), 34 deletions(-) create mode 100644 pfbudget/transform/__init__.py create mode 100644 pfbudget/transform/tagger.py diff --git a/pfbudget/core/manager.py b/pfbudget/core/manager.py index dbd0ed9..ca749f8 100644 --- a/pfbudget/core/manager.py +++ b/pfbudget/core/manager.py @@ -29,6 +29,7 @@ from pfbudget.extract.parsers import parse_data from pfbudget.extract.psd2 import PSD2Extractor from pfbudget.transform.categorizer import Categorizer from pfbudget.transform.nullifier import Nullifier +from pfbudget.transform.tagger import Tagger class Manager: @@ -107,6 +108,9 @@ class Manager: Categorizer().rules(uncategorized, categories, tags, params[0]) + rules = [rule for tag in tags for rule in tag.rules] + Tagger(rules).transform_inplace(uncategorized) + case Operation.BankMod: with self.db.session() as session: session.update(Bank, params) diff --git a/pfbudget/transform/__init__.py b/pfbudget/transform/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pfbudget/transform/categorizer.py b/pfbudget/transform/categorizer.py index 451f55c..1f59b03 100644 --- a/pfbudget/transform/categorizer.py +++ b/pfbudget/transform/categorizer.py @@ -31,7 +31,6 @@ class Categorizer: categories = [cat for cat in categories if cat.name != "null"] self._rule_based_categories(transactions, categories) - self._rule_based_tags(transactions, tags) @Timer(name="categoryrules") def _rule_based_categories( @@ -77,33 +76,3 @@ class Categorizer: for k, v in d.items(): print(f"{v}: {k}") - - @Timer(name="tagrules") - def _rule_based_tags( - self, transactions: Sequence[t.BankTransaction], tags: Sequence[t.Tag] - ): - print(f"Tagging {len(transactions)} transactions") - d = {} - for tag in [t for t in tags if len(t.rules) > 0]: - for rule in tag.rules: - # for transaction in [t for t in transactions if not t.category]: - for transaction in [ - t - for t in transactions - if tag.name not in [tag.tag for tag in t.tags] - ]: - if not rule.matches(transaction): - continue - - if not transaction.tags: - transaction.tags = {t.TransactionTag(tag.name)} - else: - transaction.tags.add(t.TransactionTag(tag.name)) - - if rule in d: - d[rule] += 1 - else: - d[rule] = 1 - - for k, v in d.items(): - print(f"{v}: {k}") diff --git a/pfbudget/transform/tagger.py b/pfbudget/transform/tagger.py new file mode 100644 index 0000000..5e2dd81 --- /dev/null +++ b/pfbudget/transform/tagger.py @@ -0,0 +1,30 @@ +from copy import deepcopy +from typing import Sequence + +from pfbudget.db.model import TagRule, Transaction, TransactionTag +from .transform import Transformer + + +class Tagger(Transformer): + def __init__(self, rules: Sequence[TagRule]): + self.rules = rules + + def transform(self, transactions: Sequence[Transaction]) -> Sequence[Transaction]: + result = deepcopy(transactions) + self.transform_inplace(result) + + return result + + def transform_inplace(self, transactions: Sequence[Transaction]) -> None: + for rule in self.rules: + for transaction in transactions: + if rule.tag in transaction.tags: + continue + + if not rule.matches(transaction): + continue + + if not transaction.tags: + transaction.tags = {TransactionTag(rule.tag)} + else: + transaction.tags.add(TransactionTag(rule.tag)) diff --git a/tests/mocks/categories.py b/tests/mocks/categories.py index f111ea3..ad5b8a6 100644 --- a/tests/mocks/categories.py +++ b/tests/mocks/categories.py @@ -1,11 +1,15 @@ from decimal import Decimal -from pfbudget.db.model import Category, CategoryRule +from pfbudget.db.model import Category, CategoryRule, Tag, TagRule category_null = Category("null", None, set()) -category_cat1 = Category( +category1 = Category( "cat#1", None, {CategoryRule(None, None, "desc#1", None, None, None, Decimal(0), "cat#1")}, ) + +tag_1 = Tag( + "tag#1", {TagRule(None, None, "desc#1", None, None, None, Decimal(0), "tag#1")} +) diff --git a/tests/test_transform.py b/tests/test_transform.py index fc01dc1..ac6d05a 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -10,9 +10,11 @@ from pfbudget.db.model import ( CategorySelector, Selector_T, TransactionCategory, + TransactionTag, ) from pfbudget.transform.categorizer import Categorizer from pfbudget.transform.nullifier import Nullifier +from pfbudget.transform.tagger import Tagger from pfbudget.transform.transform import Transformer @@ -77,6 +79,20 @@ class TestTransform: "null", CategorySelector(Selector_T.nullifier) ) + def test_tagger(self): + transactions = [ + BankTransaction(date(2023, 1, 1), "desc#1", Decimal("-10"), "Bank#1") + ] + + for t in transactions: + assert not t.category + + categorizer = Tagger(mock.tag_1.rules) + transactions = categorizer.transform(transactions) + + for t in transactions: + assert TransactionTag("tag#1") in t.tags + def test_categorize(self): transactions = [ BankTransaction(date(2023, 1, 1), "desc#1", Decimal("-10"), "Bank#1") @@ -86,7 +102,7 @@ class TestTransform: assert not t.category categorizer = Categorizer() - categorizer.rules(transactions, [mock.category_cat1], []) + categorizer.rules(transactions, [mock.category1], []) for t in transactions: assert t.category == TransactionCategory(