Merge pull request #692 from tcely/patch-4

Add & use functions from status classes

tubesync/sync/hooks.py

@@ -5,35 +5,86 @@ from common.logger import log
 from django.conf import settings
 
 
-class ProgressHookStatus:
+progress_hook = {
+    'status': dict(),
+}
+
+postprocessor_hook = {
+    'status': dict(),
+}
+
+
+class BaseStatus:
+    status_dict = dict()
+    valid = set()
+
+    @classmethod
+    def get(cls, key):
+        return cls.status_dict.get(key, None)
+
+    @classmethod
+    def valid_status(cls, status):
+        return status in cls.valid
+
+    def __init__(self, hook_status_dict=None):
+        self._status_dict = hook_status_dict or self.status_dict
+        self._registered_keys = set()
+
+    def register(self, *args):
+        additions = dict()
+        for key in args:
+            if key is not None:
+                self._registered_keys.add(key)
+                additions[key] = self
+        self._status_dict.update(additions)
+
+    def cleanup(self):
+        for key in self._registered_keys:
+            if key in self._status_dict:
+                del self._status_dict[key]
+
+class ProgressHookStatus(BaseStatus):
+    status_dict = progress_hook['status']
     valid = frozenset((
         'downloading',
         'finished',
         'error',
     ))
 
-    def __init__(self):
+    def __init__(self, *args, status=None, info_dict={}, filename=None, **kwargs):
+        super().__init__(self.status_dict)
+        self.filename = filename
+        self.info = info_dict
+        self.status = status
         self.download_progress = 0
 
-class PPHookStatus:
+    def next_progress(self):
+        if 0 == self.download_progress:
+            return 0
+        return 1 + self.download_progress
+
+class PPHookStatus(BaseStatus):
+    status_dict = postprocessor_hook['status']
     valid = frozenset((
         'started',
         'processing',
         'finished',
     ))
 
-    def __init__(self, *args, status=None, postprocessor=None, info_dict={}, **kwargs):
+    def __init__(self, *args, status=None, postprocessor=None, info_dict={}, filename=None, **kwargs):
+        super().__init__(self.status_dict)
+        self.filename = filename
         self.info = info_dict
+        self.media_name = None
         self.name = postprocessor
         self.status = status
 
-
 def yt_dlp_progress_hook(event):
-    if event['status'] not in ProgressHookStatus.valid:
+    if not ProgressHookStatus.valid_status(event['status']):
         log.warn(f'[youtube-dl] unknown progress event: {str(event)}')
         return None
 
-    name = key = 'Unknown'
+    key = None
     if 'display_id' in event['info_dict']:
         key = event['info_dict']['display_id']
     elif 'id' in event['info_dict']:
@@ -43,14 +94,17 @@ def yt_dlp_progress_hook(event):
     if 'error' == event['status']:
         log.error(f'[youtube-dl] error occured downloading: {filename}')
     elif 'downloading' == event['status']:
-        # get or create the status for key
-        status = progress_hook['status'].get(key, None)
+        # get or create the status for filename
+        status = ProgressHookStatus.get(filename)
         if status is None:
-            progress_hook['status'].update({key: ProgressHookStatus()})
+            status = ProgressHookStatus(**event)
+            status.register(key, filename, status.filename)
 
         downloaded_bytes = event.get('downloaded_bytes', 0) or 0
         total_bytes_estimate = event.get('total_bytes_estimate', 0) or 0
         total_bytes = event.get('total_bytes', 0) or total_bytes_estimate
+        fragment_index = event.get('fragment_index', 0) or 0
+        fragment_count = event.get('fragment_count', 0) or 0
         eta = event.get('_eta_str', '?').strip()
         percent_str = event.get('_percent_str', '?').strip()
         speed = event.get('_speed_str', '?').strip()
@@ -60,17 +114,21 @@ def yt_dlp_progress_hook(event):
             percent = int(float(percent_str.rstrip('%')))
         except:
             pass
-        if downloaded_bytes > 0 and total_bytes > 0:
+        if fragment_index >= 0 and fragment_count > 0:
+            percent = round(100 * fragment_index / fragment_count)
+            percent_str = f'{percent}%'
+        elif downloaded_bytes >= 0 and total_bytes > 0:
             percent = round(100 * downloaded_bytes / total_bytes)
-        if percent and (0 < percent) and (0 == percent % 5):
+        if percent and (status.next_progress() < percent) and (0 == percent % 5):
+            status.download_progress = percent
             log.info(f'[youtube-dl] downloading: {filename} - {percent_str} '
                      f'of {total} at {speed}, {eta} remaining')
-        status.download_progress = percent or 0
     elif 'finished' == event['status']:
-        # update the status for key to the finished value
-        status = progress_hook['status'].get(key, None)
+        # update the status for filename to the finished value
+        status = ProgressHookStatus.get(filename)
         if status is None:
-            progress_hook['status'].update({key: ProgressHookStatus()})
+            status = ProgressHookStatus(**event)
+            status.register(key, filename, status.filename)
         status.download_progress = 100
 
         total_size_str = event.get('_total_bytes_str', '?').strip()
@@ -78,22 +136,22 @@ def yt_dlp_progress_hook(event):
         log.info(f'[youtube-dl] finished downloading: {filename} - '
                  f'{total_size_str} in {elapsed_str}')
 
-        # clean up the status for key
-        if key in progress_hook['status']:
-            del progress_hook['status'][key]
+        status.cleanup()
 
 def yt_dlp_postprocessor_hook(event):
-    if event['status'] not in PPHookStatus.valid:
+    if not PPHookStatus.valid_status(event['status']):
         log.warn(f'[youtube-dl] unknown postprocessor event: {str(event)}')
         return None
 
     name = key = 'Unknown'
+    filename = os.path.basename(event.get('filename', '???'))
     if 'display_id' in event['info_dict']:
         key = event['info_dict']['display_id']
     elif 'id' in event['info_dict']:
         key = event['info_dict']['id']
 
-    postprocessor_hook['status'].update({key: PPHookStatus(*event)})
+    status = PPHookStatus(**event)
+    status.register(key, filename, status.filename)
 
     title = None
     if 'fulltitle' in event['info_dict']:
@@ -104,6 +162,8 @@ def yt_dlp_postprocessor_hook(event):
     if title:
         name = f'{key}: {title}'
 
+    status.media_name = name
+
     if 'started' == event['status']:
         if 'formats' in event['info_dict']:
             del event['info_dict']['formats']
@@ -112,19 +172,17 @@ def yt_dlp_postprocessor_hook(event):
         log.debug(repr(event['info_dict']))
 
     log.info(f'[{event["postprocessor"]}] {event["status"]} for: {name}')
-    if 'finished' == event['status'] and key in postprocessor_hook['status']:
-        del postprocessor_hook['status'][key]
+    if 'finished' == event['status']:
+        status.cleanup()
 
 
-progress_hook = {
+progress_hook.update({
     'class': ProgressHookStatus(),
     'function': yt_dlp_progress_hook,
-    'status': dict(),
-}
+})
 
-postprocessor_hook = {
+postprocessor_hook.update({
     'class': PPHookStatus(),
     'function': yt_dlp_postprocessor_hook,
-    'status': dict(),
-}
+})