refactor(banks): revert to having extract() handle pages

benjamin-awd · Oct 21, 2023 · 5460c83 · 5460c83
1 parent 444437f
commit 5460c83
Show file tree

Hide file tree

Showing 8 changed files with 12 additions and 14 deletions.
diff --git a/monopoly/examples/example_bank.py b/monopoly/examples/example_bank.py
@@ -18,3 +18,5 @@ class MonopolyBank(BankBase):
         statement_date_pattern=r"\d{2}\-\d{2}\-\d{4}",
         statement_date_format=r"%d-%m-%Y",
     )
+
+    pdf_config = None
diff --git a/monopoly/examples/multiple_statements.py b/monopoly/examples/multiple_statements.py
@@ -12,8 +12,7 @@ def ocbc_example():
         bank = Ocbc(
             file_path=file_path,
         )
-        pages = bank.get_pages()
-        statement = bank.extract(pages)
+        statement = bank.extract()
         transformed_df = bank.transform(statement)
         bank.load(transformed_df, statement)
 

diff --git a/monopoly/examples/single_statement.py b/monopoly/examples/single_statement.py
@@ -11,8 +11,7 @@ def example():
 
     # This runs Tesseract on the PDF and
     # extracts transactions as raw text
-    pages = bank.get_pages()
-    statement = bank.extract(pages)
+    statement = bank.extract()
 
     # Dates are converted into an ISO 8601 date format
     transformed_df = bank.transform(statement)

diff --git a/monopoly/main.py b/monopoly/main.py
@@ -38,8 +38,7 @@ def process_bank_statement(message: Message, banks: dict):
 
             with message.save(attachment) as file:
                 processor: StatementProcessor = bank_class(file_path=file)
-                pages = processor.get_pages()
-                statement = processor.extract(pages)
+                statement = processor.extract()
                 transformed_df = processor.transform(statement)
                 processor.load(transformed_df, statement, upload_to_cloud=True)
 

diff --git a/monopoly/processor.py b/monopoly/processor.py
@@ -6,7 +6,7 @@
 
 from monopoly.config import PdfConfig, StatementConfig, settings
 from monopoly.constants import StatementFields
-from monopoly.pdf import PdfPage, PdfParser
+from monopoly.pdf import PdfParser
 from monopoly.statement import Statement
 from monopoly.storage import upload_to_cloud_storage, write_to_csv
 
@@ -24,7 +24,9 @@ def __init__(
 
         super().__init__(file_path=self.file_path, config=pdf_config)
 
-    def extract(self, pages: list[PdfPage]) -> Statement:
+    def extract(self) -> Statement:
+        parser = PdfParser(self.file_path, self.pdf_config)
+        pages = parser.get_pages()
         statement = Statement(pages, self.statement_config)
 
         if not statement.transactions:

diff --git a/tests/integration/banks/citibank/test_citibank_extract.py b/tests/integration/banks/citibank/test_citibank_extract.py
@@ -6,8 +6,7 @@
 
 
 def test_citibank_extract_unprotected_pdf(citibank: Citibank):
-    pages = citibank.get_pages()
-    raw_df = citibank.extract(pages).df
+    raw_df = citibank.extract().df
     expected_df = pd.read_csv("tests/integration/fixtures/citibank/expected.csv")
 
     assert_frame_equal(raw_df, expected_df)

diff --git a/tests/integration/banks/hsbc/test_hsbc_extract.py b/tests/integration/banks/hsbc/test_hsbc_extract.py
@@ -6,8 +6,7 @@
 
 
 def test_hsbc_extract_unprotected_pdf(hsbc: Hsbc):
-    pages = hsbc.get_pages()
-    raw_df = hsbc.extract(pages).df
+    raw_df = hsbc.extract().df
     expected_df = pd.read_csv("tests/integration/fixtures/hsbc/expected.csv")
 
     assert_frame_equal(raw_df, expected_df)

diff --git a/tests/integration/banks/ocbc/test_ocbc_extract.py b/tests/integration/banks/ocbc/test_ocbc_extract.py
@@ -6,8 +6,7 @@
 
 
 def test_ocbc_extract_unprotected_pdf(ocbc: Ocbc):
-    pages = ocbc.get_pages()
-    raw_df = ocbc.extract(pages).df
+    raw_df = ocbc.extract().df
 
     expected_df = pd.read_csv("tests/integration/fixtures/ocbc/expected.csv")