chore: improved document structure

src/weby_pilot/__main__.py

@@ -3,12 +3,12 @@
 
 from .bpi import BpiAPI
 
-# BpiAPI.download_report(report_indexes=range(0, 8))
+print(BpiAPI().download_account_report(report_indexes=range(0, 2)))
 
-# BpiAPI.download_report(section="Extrato Investimento", report_indexes=range(0, 8))
+# BpiAPI().download_report(section="Extrato Investimento", report_indexes=range(0, 8))
 
-# BpiAPI.download_invoice()
+# BpiAPI().download_invoice()
 
-print(BpiAPI().get_balance())
+# print(BpiAPI().get_balance())
 
-# BpiAPI.download_card_report(report_indexes=range(0, 8))
+# BpiAPI().download_card_report(report_indexes=range(0, 8))

src/weby_pilot/base.py

@@ -1,7 +1,8 @@
 #!/usr/bin/python
 # -*- coding: utf-8 -*-
 
-from typing import Any, Generator, List, Literal
+from io import BytesIO
+from typing import IO, Any, Generator, List, Literal
 from uuid import uuid4
 from time import sleep
 from contextlib import contextmanager
@@ -29,6 +30,7 @@ class WebyAPI:
     _wait: WebDriverWait[Chrome] | None = None
     _downloads_dir: str = abspath("downloads")
     _temp_dir: str = abspath("temp")
+    _last_path: str | None = None
 
     @classmethod
     def build_options(cls):
@@ -131,7 +133,10 @@ def wait_download(
 
                 rename(src, dst)
 
+            self._last_path = dst
+
             # if we reach this point, then the download is completed
+            # we can return the destination file path
             return dst
 
         raise Exception(f"Download not completed after {timeout} seconds")
@@ -191,3 +196,25 @@ def wait(self) -> WebDriverWait[Chrome]:
     def _cleanup_temp(self):
         for filename in listdir(self._temp_dir):
             remove(f"{self._temp_dir}/{filename}")
+
+    def _last_download(self) -> IO:
+        return open(self._last_download_path, "rb")
+
+    def _last_download_buffer(self) -> IO:
+        file = self._last_download()
+        try:
+            content = file.read()
+        finally:
+            file.close()
+        buffer = BytesIO(content)
+        buffer.write(content)
+        buffer.seek(0)
+        return buffer
+
+    @property
+    def _last_download_path(self) -> str:
+        if self._last_path is None:
+            raise Exception("No file downloaded")
+        if not exists(self._last_path):
+            raise Exception(f"File {self._last_download} does not exist")
+        return self._last_path

src/weby_pilot/bpi.py

@@ -1,8 +1,10 @@
 #!/usr/bin/python
 # -*- coding: utf-8 -*-
 
+from enum import Enum
 from os import environ
-from typing import Literal, Sequence, Tuple, cast
+from datetime import datetime
+from typing import IO, Literal, Sequence, Tuple, cast
 from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
 
@@ -81,7 +83,6 @@ def build_login(cls) -> Tuple[str, str]:
         return username, password
 
     def get_balance(self) -> str:
-        cls = self.__class__
         with self.driver_ctx():
             self.login()
             self.select_section("Consultas")
@@ -93,28 +94,61 @@ def download_invoice(
         date_range: SelectDateRange = "Último Ano",
         document_type: DocumentType = "Facturas",
         invoice_indexes: Sequence[int] = (0,),
-    ):
+    ) -> Sequence["BpiDocument"]:
+        docs: list[BpiDocument] = []
         with self.driver_ctx():
-            self.login(*cls.build_login())
+            self.login()
             self.select_section("Consultas")
             self.select_side_menu("Avisos/Faturas/Notas Crédito e Débito")
             self.select_filters(date_range=date_range, document_type=document_type)
             for invoice_index in invoice_indexes:
                 self.click_extract(row_index=invoice_index)
+                docs.append(
+                    BpiDocument(
+                        BpiDocumentType.from_section(document_type),
+                        self._last_download_path,
+                        self._last_download_buffer(),
+                        date=datetime.strptime(
+                            self._last_download_path[-14:-4], "%Y-%m-%d"
+                        ),
+                    )
+                )
+        return docs
 
     def download_report(
         self,
         section: ReportSections = "Extrato Conta",
         report_indexes: Sequence[int] = (0,),
-    ):
+    ) -> Sequence["BpiDocument"]:
+        docs: list[BpiDocument] = []
         with self.driver_ctx():
             self.login()
             self.select_section("Consultas")
             self.select_side_menu(cast(BpiSideSections, section))
             for report_index in report_indexes:
                 self.click_extract(row_index=report_index)
+                docs.append(
+                    BpiDocument(
+                        BpiDocumentType.from_section(section),
+                        self._last_download_path,
+                        self._last_download_buffer(),
+                        date=datetime.strptime(
+                            self._last_download_path[-14:-4], "%Y-%m-%d"
+                        ),
+                    )
+                )
+        return docs
 
-    def download_investing_report(self, report_indexes: Sequence[int] = (0,)):
+    def download_account_report(
+        self, report_indexes: Sequence[int] = (0,)
+    ) -> Sequence["BpiDocument"]:
+        return self.download_report(
+            section="Extrato Conta", report_indexes=report_indexes
+        )
+
+    def download_investing_report(
+        self, report_indexes: Sequence[int] = (0,)
+    ) -> Sequence["BpiDocument"]:
         return self.download_report(
             section="Extrato Investimento", report_indexes=report_indexes
         )
@@ -176,3 +210,47 @@ def click_card_account(self, row_index=0, wait_download: bool = True):
             By.XPATH, f"//a[contains(@class, 'Text_NoWrap')][{row_index + 1}]"
         )
         card_account.click()
+
+
+class BpiDocumentType(Enum):
+    ACCOUNT_EXTRACT = 1
+    CARD_EXTRACT = 2
+    INVESTMENT_EXTRACT = 3
+    INVOICE = 4
+    UNKNOWN = 100
+
+    @classmethod
+    def from_section(cls, section: str) -> "BpiDocumentType":
+        if section == "Extrato Conta":
+            return BpiDocumentType.ACCOUNT_EXTRACT
+        if section == "Extrato Cartões":
+            return BpiDocumentType.CARD_EXTRACT
+        if section == "Extrato Investimento":
+            return BpiDocumentType.INVESTMENT_EXTRACT
+        if section == "Avisos/Faturas/Notas Crédito e Débito":
+            return BpiDocumentType.INVOICE
+        if section == "Faturas":
+            return BpiDocumentType.INVOICE
+        return BpiDocumentType.UNKNOWN
+
+    def __repr__(self):
+        return self.name.lower().replace("_", " ").title()
+
+
+class BpiDocument:
+
+    type: BpiDocumentType
+    name: str
+    buffer: IO
+    date: datetime | None
+
+    def __init__(
+        self, type: BpiDocumentType, name: str, buffer: IO, date: datetime | None = None
+    ):
+        self.type = type
+        self.name = name
+        self.buffer = buffer
+        self.date = date
+
+    def __repr__(self):
+        return f"BpiDocument(type={self.type}, name={self.name}, date={self.date})"