diff --git a/pillarsdk/files.py b/pillarsdk/files.py index b3250d5..fe34e75 100755 --- a/pillarsdk/files.py +++ b/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() diff --git a/pillarsdk/nodes.py b/pillarsdk/nodes.py index e75675e..cc781f2 100755 --- a/pillarsdk/nodes.py +++ b/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