Merge pull request #73 from SchrodingersGat/file-downloader

Adds API function for downloading file

Author: SchrodingersGat

inventree/api.py

@@ -35,7 +35,8 @@ def __init__(self, base_url, **kwargs):
         """ Initialize class with initial parameters
 
         Args:
-            base_url - Base API URL
+            base_url - Base URL for the InvenTree server, including port (if required)
+                       e.g. "http://inventree.server.com:8000"
             
         kwargs:
             username - Login username
@@ -45,18 +46,17 @@ def __init__(self, base_url, **kwargs):
             verbose - Print extra debug messages (default = False)
         """
 
-        if not base_url.endswith('/'):
-            base_url += '/'
-
-        # Server address *must* end with /api/
-        if not base_url.endswith('/api/'):
-            base_url = os.path.join(base_url, 'api')
+        # Strip out trailing "/api/" (if provided)
+        if base_url.endswith("/api/"):
+            base_url = base_url[:-5]
 
         if not base_url.endswith('/'):
             base_url += '/'
 
         self.base_url = base_url
 
+        self.api_url = os.path.join(self.base_url, 'api/')
+
         logger.info("Connecting to server: " + str(self.base_url))
 
         self.username = kwargs.get('username', None)
@@ -78,7 +78,7 @@ def __init__(self, base_url, **kwargs):
 
     def clean_url(self, url):
 
-        url = os.path.join(self.base_url, url)
+        url = os.path.join(self.api_url, url)
 
         if not url.endswith('/'):
             url += '/'
@@ -97,7 +97,7 @@ def testServer(self):
         logger.info("Checking InvenTree server connection...")
 
         try:
-            response = requests.get(self.base_url)
+            response = requests.get(self.api_url)
         except requests.exceptions.ConnectionError:
             logger.error("Server connection refused - check server address")
             return False
@@ -144,7 +144,7 @@ def requestToken(self):
         logger.info("Requesting auth token from server...")
 
         # Request an auth token from the server
-        token_url = os.path.join(self.base_url, 'user/token/')
+        token_url = os.path.join(self.api_url, 'user/token/')
 
         reply = requests.get(token_url, auth=self.auth)
 
@@ -176,7 +176,7 @@ def request(self, url, **kwargs):
         if url.startswith('/'):
             url = url[1:]
 
-        api_url = os.path.join(self.base_url, url)
+        api_url = os.path.join(self.api_url, url)
 
         if not api_url.endswith('/'):
             api_url += '/'
@@ -426,3 +426,55 @@ def get(self, url, **kwargs):
             return None
 
         return data
+
+    def downloadFile(self, url, destination):
+        """
+        Download a file from the InvenTree server.
+
+        - If the "destination" is a directory, use the filename of the remote URL
+        """
+
+        # Check that the provided URL is "absolute"
+        if not url.startswith(self.base_url):
+
+            if url.startswith('/'):
+                url = url[1:]
+
+            url = os.path.join(self.base_url, url)
+
+        if os.path.exists(destination) and os.path.isdir(destination):
+
+            destination = os.path.join(
+                destination,
+                os.path.basename(url)
+            )
+
+        destination = os.path.abspath(destination)
+
+        headers = {
+            'AUTHORIZATION': f"Token {self.token}"
+        }
+
+        with requests.get(url, stream=True, headers=headers) as request:
+
+            if not request.status_code == 200:
+                logger.error(
+                    f"Error downloading file '{url}': Server returned status {request.status_code}"
+                )
+                return False
+
+            headers = request.headers
+
+            if 'text/html' in headers['Content-Type']:
+                logger.error(
+                    f"Error downloading file '{url}': Server return invalid response (text/html)"
+                )
+                return False
+
+            with open(destination, 'wb') as f:
+
+                for chunk in request.iter_content(chunk_size=16 * 1024):
+                    f.write(chunk)
+
+        logger.info(f"Downloaded '{url}' to '{destination}'")
+        return True

inventree/base.py

@@ -226,6 +226,13 @@ def upload(cls, api, filename, comment, **kwargs):
         else:
             logger.warning("File upload failed")
 
+    def download(self, destination):
+        """
+        Download the attachment file to the specified location
+        """
+
+        return self._api.downloadFile(self.attachment, destination)
+
 
 class Currency(InventreeObject):
     """ Class representing the Currency database model """

inventree/company.py

@@ -1,9 +1,15 @@
 # -*- coding: utf-8 -*-
 
+import os
+import logging
+
 import inventree.base
 import inventree.order
 
 
+logger = logging.getLogger('inventree')
+
+
 class Company(inventree.base.InventreeObject):
     """ Class representing the Company database model """
 
@@ -57,6 +63,42 @@ def createSalesOrder(self, **kwargs):
             data=kwargs
         )
 
+    def uploadImage(self, image):
+        """
+        Upload an image file against this Company
+        
+        Returns the HTTP response object
+        """
+
+        files = {}
+
+        if image:
+            if os.path.exists(image):
+                f = os.path.basename(image)
+                fo = open(image, 'rb')
+                files['image'] = (f, fo)
+            else:
+                logger.error(f"Image '{image}' does not exist")
+                return None
+
+        response = self.save(
+            data={},
+            files=files,
+        )
+
+        return response
+
+    def downloadImage(self, destination):
+        """
+        Download the image for this Company, to the specified destination
+        """
+
+        if self.image:
+            return self._api.downloadFile(self.image, destination)
+        else:
+            logger.error(f"Company '{self.name}' does not have an associated image")
+            return False
+
 
 class SupplierPart(inventree.base.InventreeObject):
     """ Class representing the SupplierPart database model """

inventree/part.py

@@ -10,6 +10,9 @@
 import inventree.build
 
 
+logger = logging.getLogger('inventree')
+
+
 class PartCategory(inventree.base.InventreeObject):
     """ Class representing the PartCategory database model """
 
@@ -98,7 +101,7 @@ def uploadImage(self, image):
                 fo = open(image, 'rb')
                 files['image'] = (f, fo)
             else:
-                logging.error("File does not exist: '{f}'".format(f=image))
+                logger.error("File does not exist: '{f}'".format(f=image))
                 return None
 
         response = self.save(
@@ -108,6 +111,17 @@ def uploadImage(self, image):
 
         return response
 
+    def downloadImage(self, destination):
+        """
+        Download the image for this Part, to the specified destination
+        """
+
+        if self.image:
+            return self._api.downloadFile(self.image, destination)
+        else:
+            logger.error(f"Part '{self.name}' does not have an associated image")
+            return False
+
 
 class PartAttachment(inventree.base.Attachment):
     """ Class representing a file attachment for a Part """
@@ -139,7 +153,7 @@ def upload_attachment(cls, api, part, **kwargs):
                 fo = open(attachment, 'rb')
                 files['attachment'] = (f, fo)
             else:
-                logging.error("File does not exist: '{f}'".format(f=attachment))
+                logger.error("File does not exist: '{f}'".format(f=attachment))
 
         comment = kwargs.get('comment', '')
 
@@ -151,10 +165,10 @@ def upload_attachment(cls, api, part, **kwargs):
 
         # Send the data to the server
         if api.post(cls.URL, data, files=files):
-            logging.info("Uploaded attachment: '{f}'".format(f=attachment))
+            logger.info("Uploaded attachment: '{f}'".format(f=attachment))
             ret = True
         else:
-            logging.warning("Attachment upload failed")
+            logger.warning("Attachment upload failed")
             ret = False
 
         return ret

test/test_api.py

@@ -46,6 +46,13 @@ def test_read_parts(self):
 
         self.assertEqual(len(parts), 0)
 
+    def test_file_download(self):
+        """
+        Attemtping to download a file while unauthenticated should return False
+        """
+
+        self.assertFalse(self.api.downloadFile('/media/part/files/1/test.pdf', 'test.pdf'))
+
 
 class InvenTreeTestCase(unittest.TestCase):
 

test/test_part.py

@@ -214,6 +214,9 @@ def test_image_upload(self):
         # Grab the first part
         p = part.Part.list(self.api)[0]
 
+        # Ensure the part does *not* have an image associated with it
+        p.save(data={'image': None})
+
         # Create a dummy file (not an image)
         with open('dummy_image.jpg', 'w') as dummy_file:
             dummy_file.write("hello world")
@@ -231,3 +234,13 @@ def test_image_upload(self):
         self.assertIsNotNone(response)
         self.assertIsNotNone(p['image'])
         self.assertIn('dummy_image', p['image'])
+
+        # Re-download the image
+        self.assertTrue(p.downloadImage("downloaded_image.png"))
+
+        self.assertTrue(os.path.exists("downloaded_image.png"))
+
+        with Image.open("downloaded_image.png") as img2:
+
+            self.assertEqual(img2.width, 128)
+            self.assertEqual(img2.height, 128)