ExecuteExternalStorage.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. import os
  2. import uuid
  3. from flask import abort
  4. from flask import make_response
  5. from flask import request
  6. from flask.views import View
  7. from pycs import db
  8. from pycs.database.Project import Project
  9. from pycs.frontend.notifications.NotificationManager import NotificationManager
  10. from pycs.jobs.JobGroupBusyException import JobGroupBusyException
  11. from pycs.jobs.JobRunner import JobRunner
  12. from pycs.util.FileOperations import file_info
  13. class ExecuteExternalStorage(View):
  14. """
  15. find media files stored in a projects data_folder
  16. """
  17. # pylint: disable=arguments-differ
  18. methods = ['POST']
  19. def __init__(self, nm: NotificationManager, jobs: JobRunner):
  20. # pylint: disable=invalid-name
  21. self.nm = nm
  22. self.jobs = jobs
  23. def dispatch_request(self, project_id: int):
  24. # extract request data
  25. data = request.get_json(force=True)
  26. if not data.get('execute', False):
  27. return abort(400)
  28. # find project
  29. project = Project.get_or_404(project_id)
  30. if not project.external_data:
  31. return abort(400, "External data is not set!")
  32. # execute label provider and add labels to project
  33. try:
  34. ExecuteExternalStorage.find_media_files(self.nm, self.jobs, project)
  35. except JobGroupBusyException:
  36. return abort(400, "Job is already running!")
  37. return make_response()
  38. @staticmethod
  39. def find_media_files(nm: NotificationManager, jobs: JobRunner, project: Project):
  40. """
  41. start a job that finds media files in the projects data_folder and adds them to the
  42. database afterwards
  43. :param nm: notification manager object
  44. :param jobs: job runner object
  45. :param project: project
  46. :return:
  47. """
  48. # pylint: disable=invalid-name
  49. # find lists the given data folder and prepares item dictionaries
  50. def find(data_folder):
  51. files = os.listdir(data_folder)
  52. length = len(files)
  53. elements = []
  54. current = 0
  55. for file_name in files:
  56. file_path = os.path.join(data_folder, file_name)
  57. if not os.path.isfile(file_path):
  58. continue
  59. file_name, file_extension = os.path.splitext(file_name)
  60. file_size = os.path.getsize(file_path)
  61. try:
  62. ftype, frames, fps = file_info(data_folder, file_name, file_extension)
  63. except ValueError:
  64. continue
  65. file_attrs = dict(
  66. uuid=str(uuid.uuid1()),
  67. file_type=ftype,
  68. name=file_name,
  69. extension=file_extension,
  70. size=file_size,
  71. frames=frames,
  72. fps=fps)
  73. elements.append(file_attrs)
  74. current += 1
  75. if len(elements) >= 200:
  76. yield elements, current, length
  77. elements = []
  78. if len(elements) > 0:
  79. yield elements, current, length
  80. # progress inserts elements into the database and fires events
  81. def progress(elements, current, length):
  82. with db.session.begin_nested():
  83. for file_attrs in elements:
  84. file, is_new = project.add_file(commit=False, **file_attrs)
  85. if is_new:
  86. nm.create_file(file)
  87. return current / length
  88. # run job with given functions
  89. jobs.run(project,
  90. 'Find Media Files',
  91. project.name,
  92. f'{project.id}/find-files',
  93. find, project.data_folder,
  94. progress=progress)