PyCS/pycs/frontend/FileProvider.py

from os import path, mkdir, getcwd
from time import time
from uuid import uuid1

from flask import Flask, make_response, send_from_directory, request
from werkzeug import formparser

from pycs.ApplicationStatus import ApplicationStatus
from pycs.util.GenericWrapper import GenericWrapper
from pycs.util.ProgressFileWriter import ProgressFileWriter


class FileProvider(Flask):
    def __init__(self, app_status: ApplicationStatus, cors=False):
        super().__init__(__name__)

        # add download handler
        @self.route('/media/<identifier>', methods=['GET'])
        def media(identifier):
            # get current project
            opened_projects = list(filter(lambda x: x['status'] == 'open', app_status['projects']))
            if len(opened_projects) == 0:
                return make_response('no open project available', 500)

            current_project = opened_projects[0]
            file_directory = path.join(getcwd(), 'projects', current_project['id'], 'data')

            # get object
            data_objects = list(filter(lambda x: x['id'] == identifier, current_project['data']))
            if len(data_objects) == 0:
                return make_response('data object not avilable', 500)

            target_object = data_objects[0]

            # return data
            file_name = target_object['id'] + target_object['extension']
            return send_from_directory(file_directory, file_name)

        # add upload handler
        @self.route('/upload', methods=['POST'])
        def upload():
            file_uuid = str(uuid1())

            # get current project path
            opened_projects = list(filter(lambda x: x['status'] == 'open', app_status['projects']))
            if len(opened_projects) == 0:
                return make_response('no open project available', 500)

            current_project = opened_projects[0]
            upload_path = path.join('projects', current_project['id'], 'data')

            job = GenericWrapper()
            file_name = GenericWrapper()
            file_extension = GenericWrapper()
            file_size = GenericWrapper(0)

            # save upload to file
            def custom_stream_factory(total_content_length, filename, content_type, content_length=None):
                file_name.value, file_extension.value = path.splitext(filename)
                file_path = path.join(upload_path, f'{file_uuid}{file_extension.value}')

                # add job to app status
                job.value = app_status['jobs'].append({
                    'id': file_uuid,
                    'type': 'upload',
                    'progress': 0,
                    'filename': filename,
                    'created': int(time()),
                    'finished': None
                })

                # create upload path if not exists
                if not path.exists(upload_path):
                    mkdir(upload_path)

                # define progress callback
                length = content_length if content_length is not None and content_length != 0 else total_content_length

                def callback(progress):
                    file_size.value += progress
                    relative = progress / length

                    if relative - job.value['progress'] > 0.02:
                        job.value['progress'] = relative

                # open file handler
                return ProgressFileWriter(file_path, 'wb', callback)

            stream, form, files = formparser.parse_form_data(request.environ, stream_factory=custom_stream_factory)

            if 'file' not in files.keys():
                return make_response('no file uploaded', 500)

            # set progress to 1 after upload is done
            job = job.value

            job['progress'] = 1
            job['finished'] = int(time())

            # add to project files
            if 'data' not in current_project:
                current_project['data'] = []

            current_project['data'].append({
                'id': file_uuid,
                'name': file_name.value,
                'extension': file_extension.value,
                'size': file_size.value,
                'created': job['created']
            })

            # create response and allow cors if needed
            response = make_response()
            if cors:
                response.headers['Access-Control-Allow-Origin'] = '*'

            return response