From 96fd0a59acb4b9a618e452b06c9a7d8acf21e097 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Magalh=C3=A3es?= <joamag@gmail.com>
Date: Tue, 1 Oct 2024 11:17:44 +0100
Subject: [PATCH] chore: more file type

---
 src/weby_pilot/__init__.py |  1 +
 src/weby_pilot/bpi.py      | 22 ++++++++++++++++++++--
 src/weby_pilot/common.py   | 14 ++++++++++++++
 3 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 src/weby_pilot/common.py

diff --git a/src/weby_pilot/__init__.py b/src/weby_pilot/__init__.py
index 70ee397..10c4061 100644
--- a/src/weby_pilot/__init__.py
+++ b/src/weby_pilot/__init__.py
@@ -4,3 +4,4 @@
 from .base import WebyAPI
 from .bpi import BpiAPI
 from .big import BigAPI
+from .common import FileType
diff --git a/src/weby_pilot/bpi.py b/src/weby_pilot/bpi.py
index 7215f00..8fd80d6 100644
--- a/src/weby_pilot/bpi.py
+++ b/src/weby_pilot/bpi.py
@@ -10,6 +10,7 @@
 from selenium.webdriver.common.keys import Keys
 
 from .base import WebyAPI
+from .common import FileType
 
 BpiSections = Literal[
     "Consultas",
@@ -107,6 +108,7 @@ def download_invoice(
                         BpiDocumentType.from_section(document_type),
                         basename(self._last_download_path),
                         self._last_download_buffer(),
+                        file_type=FileType.PDF,
                         account=self.username,
                         date=datetime.strptime(
                             self._last_download_path[-14:-4], "%Y-%m-%d"
@@ -132,6 +134,7 @@ def download_report(
                         BpiDocumentType.from_section(section),
                         basename(self._last_download_path),
                         self._last_download_buffer(),
+                        file_type=FileType.PDF,
                         account=self.username,
                         date=datetime.strptime(
                             self._last_download_path[-14:-4], "%Y-%m-%d"
@@ -242,7 +245,8 @@ class BpiDocument:
 
     type: BpiDocumentType
     name: str
-    buffer: IO
+    buffer: IO[bytes]
+    file_type: FileType
     account: str | None
     date: datetime | None
 
@@ -250,15 +254,29 @@ def __init__(
         self,
         type: BpiDocumentType,
         name: str,
-        buffer: IO,
+        buffer: IO[bytes],
+        file_type: FileType,
         account: str | None = None,
         date: datetime | None = None,
     ):
         self.type = type
         self.name = name
         self.buffer = buffer
+        self.file_type = file_type
         self.account = account
         self.date = date
 
     def __repr__(self):
         return f"BpiDocument(type={self.type}, name={self.name}, account={self.account} date={self.date})"
+
+    @property
+    def year(self) -> Tuple[int, int]:
+        if self.date is None:
+            raise Exception("Date is not set")
+        return self.date.year
+
+    @property
+    def month_filename(self) -> str:
+        if self.date is None:
+            raise Exception("Date is not set")
+        return f"{self.date.strftime('%m')}.{self.file_type.extension}"
diff --git a/src/weby_pilot/common.py b/src/weby_pilot/common.py
new file mode 100644
index 0000000..b094f50
--- /dev/null
+++ b/src/weby_pilot/common.py
@@ -0,0 +1,14 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+from enum import Enum
+
+
+class FileType(Enum):
+    PDF = "application/pdf"
+    CSV = "text/csv"
+    TXT = "text/plain"
+
+    @property
+    def extension(self) -> str:
+        return self.name.lower()