Split Node.create_asset_from_file and added direct uploading from file obj

Node.create_asset_from_file() is split such that the actual file upload is
handled by File.upload_to_project(). This now also supports directly
uploading from a file object, which could be an io.BytesIO() object to
upload from RAM.

pillarsdk/files.py

@@ -1,4 +1,5 @@
 import os.path
+import logging
 
 from .resource import List
 from .resource import Find
@@ -11,6 +12,7 @@ from .resource import Replace
 from . import utils
 
 THUMBNAIL_SIZES = 'sbtmlh'
+log = logging.getLogger(__name__)
 
 
 class File(List, Find, Create, Post, Update, Delete, Replace):
@@ -114,3 +116,39 @@ class File(List, Find, Create, Post, Update, Delete, Replace):
         utils.download_to_file(thumb_link, thumb_path)
 
         return thumb_path
+
+    @classmethod
+    def upload_to_project(cls, project_id,
+                          mimetype,
+                          filename,
+                          fileobj=None,
+                          api=None):
+        """Uploads a file to the project storage space.
+
+        :param project_id: the project ID
+        :param mimetype: MIME type of the file, such as "image/jpeg"
+        :param filename: path of the file to upload. Must be readable when fileobj
+            is not given.
+        :param fileobj: file object to read the file from. If None, it is read
+            by opening 'filename'. The file object will not be closed after uploading.
+
+        :returns: the upload response as a dict {'status': 'ok', 'file_id': 'some-id'}
+        :rtype: dict
+        """
+
+        # Select which file object to upload.
+        if fileobj is None:
+            infile = open(filename, mode='rb')
+        else:
+            infile = fileobj
+            log.debug('Uploading directly from file object %r', fileobj)
+
+        assert infile is not None
+
+        # Perform the upload.
+        try:
+            return api.post('storage/stream/%s' % project_id,
+                            files={'file': (os.path.basename(filename), infile, mimetype)})
+        finally:
+            if fileobj is None:
+                infile.close()

pillarsdk/nodes.py

@@ -77,8 +77,10 @@ class Node(List, Find, Create, Post, Update, Delete, Replace):
 
     @classmethod
     def create_asset_from_file(cls, project_id, parent_node_id, asset_type, filename,
+                               mimetype=None,
                                always_create_new_node=False,
                                extra_where=None,
+                               fileobj=None,
                                api=None):
         """Uploads the file to the Cloud and creates an asset node.
 
@@ -90,24 +92,33 @@ class Node(List, Find, Create, Post, Update, Delete, Replace):
         :param parent_node_id: node ID to attach this asset node to. Can be None.
         :param asset_type: 'image', 'file', 'video', etc.
         :param filename: path of the file to upload. Must be readable.
+        :param mimetype: MIME type of the file, such as "image/jpeg". If
+            None, it will be guessed from the filename.
         :param always_create_new_node: when True, a new node is always created,
             possibly with the same name & parent as an existing one.
         :param extra_where: dict of properties to use, in addition to project, node_type
             and name, to find any existing node. Use this to restrict the
             nodes that may be re-used to attach this file to.
+        :param fileobj: file object to read the file from. If None, it is read
+            by opening 'filename'. The file object will not be closed after uploading.
         :returns: the updated/created node
         :rtype: Node
         """
 
         api = api or Api.Default()
 
-        # Upload the file.
-        with open(filename, mode='rb') as infile:
-            file_upload_resp = api.post('storage/stream/%s' % project_id,
-                                        files={'file': infile})
-        if file_upload_resp['status'] != 'ok':
+        # Guess mime type from filename.
+        if not mimetype:
+            mimetype = cls._guess_mimetype(filename)
+
+        from .files import File
+
+        # Upload the file to project storage.
+        file_upload_resp = File.upload_to_project(project_id, mimetype, filename, fileobj, api=api)
+        file_upload_status = file_upload_resp.get('_status') or file_upload_resp.get('status')
+        if file_upload_status != 'ok':
             raise ValueError('Received bad status %s from Pillar: %s' %
-                             (file_upload_resp['status'], json.dumps(file_upload_resp)))
+                             (file_upload_status, json.dumps(file_upload_resp)))
         file_id = file_upload_resp['file_id']
 
         # Create or update the node.
@@ -133,14 +144,26 @@ class Node(List, Find, Create, Post, Update, Delete, Replace):
                 return existing_node
 
         basic_properties.update({
-                'properties': {'content_type': asset_type,
-                               'file': file_id},
-            })
+            'properties': {'content_type': asset_type,
+                           'file': file_id},
+        })
         node = cls(basic_properties)
         node.create(api=api)
 
         return node
 
+    @classmethod
+    def _guess_mimetype(cls, filename):
+        """Guesses the MIME type from the filename.
+
+        :return: the MIME type
+        :rtype: str
+        """
+
+        import mimetypes
+        mimetype, _ = mimetypes.guess_type(filename, strict=False)
+        return mimetype
+
 
 class NodeType(List, Find, Create, Post, Delete):
     """NodeType class wrapping the REST node_types endpoint