Improve pdf download

python_ifirma/core.py

@@ -1,13 +1,13 @@
 # -*- coding: utf-8 -*-
 
-import six
 import json
 import datetime
-from time import strftime
+import requests
 
+from time import strftime
+from collections import defaultdict
 from python_ifirma.exceptions import PythonIfirmaExceptionFactory
 from python_ifirma.helpers import Helpers
-import requests
 
 
 from urllib.parse import urljoin, quote
@@ -176,7 +176,7 @@ class iFirmaAPI():
     __user_key_name = "abonent"
     keys = {
         __invoice_key_name: None,
-        __cost_key_name : None,
+        __cost_key_name: None,
         __user_key_name: None
     }
 
@@ -202,9 +202,10 @@ class iFirmaAPI():
             "Authentication": self.__get_auth_header(request_hash_text,
                                                      key_name)
         }
-        response = requests.get(url, headers=headers, params=params)
-        return json.loads(response.content.decode('utf-8'))
+        return requests.get(url, headers=headers, params=params)
 
+    def __get_json(self, url_path, key_name, params={}):
+        return json.loads(self.__get(url_path, key_name, params).content.decode('utf-8'))
 
     def __request(self, method, url_path, key_name, request_data={}):
         """Send request with given method to IFirma API"""
@@ -227,21 +228,20 @@ class iFirmaAPI():
         response_code = real_response_content.get("Kod", -1)
 
         if response_code not in (0,
-                                 202): # 'Numer faktury' nie jest unikalne
+                                 202):  # 'Numer faktury' nie jest unikalne
             raise PythonIfirmaExceptionFactory.throw_exception_by_code(response_code)
 
         return response_dict
 
     def find_partner(self, keyword):
-        return self.__get(f"kontrahenci/{quote(keyword)}.json",
-                          self.__invoice_key_name)
+        return self.__get_json(f"kontrahenci/{quote(keyword)}.json",
+                               self.__invoice_key_name)
 
     # Invoice generation
     def __create_invoice_and_return_id(self, invoice, url):
         response_dict = self.__request("POST", url, self.__invoice_key_name,
-                                           invoice.get_request_data())
+                                       invoice.get_request_data())
 
-        breakpoint()
         if response_dict["response"].get("Identyfikator"):
             invoice_id = response_dict["response"]["Identyfikator"]
             return invoice_id
@@ -249,8 +249,8 @@ class iFirmaAPI():
             return None
 
     def __get_invoice_number(self, invoice_id):
-        rj = self.__get(f"fakturakraj/{invoice_id}.json",
-                        self.__user_key_name)
+        rj = self.__get_json(f"fakturakraj/{invoice_id}.json",
+                             self.__user_key_name)
 
         if "Kod" in rj["response"] and rj["response"]["Kod"] != 200:
             return None
@@ -268,8 +268,8 @@ class iFirmaAPI():
         Return currently open accounting year and month
         """
 
-        response = self.__get("abonent/miesiacksiegowy.json",
-                              self.__user_key_name)
+        response = self.__get_json("abonent/miesiacksiegowy.json",
+                                   self.__user_key_name)
         return (response['response']['RokKsiegowy'],
                 response['response']["MiesiacKsiegowy"])
 
@@ -308,27 +308,13 @@ class iFirmaAPI():
         if date_to:
             params.update(dataDo=date_to)
 
-        return self.__get("faktury.json", self.__invoice_key_name,
-                          params=params)
+        return self.__get_json("faktury.json", self.__invoice_key_name,
+                               params=params)
 
-    def get_invoice_pdf(self, invoice_id):
-        url = f"https://www.ifirma.pl/iapi/fakturakraj/{invoice_id}.pdf"
-        return self.__download_pdf(url)
-
-    def __download_pdf(self, url):
-        request_hash_text = "{}{}{}".format(
-            url,
-            self.__username,
-            self.__invoice_key_name,
-        )
-        headers = {
-            "Accept": "application/pdf",
-            "Authentication": self.__get_auth_header(request_hash_text)
-        }
-        resp = requests.get(url, headers=headers)
-
-        content = resp.content
-        file_obj = six.BytesIO()
-        file_obj.write(content)
-        file_obj.seek(0)
-        return file_obj
+    # Beware: IFIRMA returns errors printed into the pdf...
+    def get_invoice_pdf(self, invoice_id, kind=""):
+        kinds = defaultdict(lambda: "fakturakraj",
+                            {"prz_eksport_dost_uslug_nie_ue": "fakturaeksportuslug",
+                             "prz_eksport_dost_uslug_ue": "fakturaeksportuslugue"})
+        return self.__get(f"{kinds[kind]}/{invoice_id}.pdf",
+                          self.__invoice_key_name)