Fix likes count

1 files changed, 42 insertions, 21 deletions

diff --git a/youtube/yt_data_extract/watch_extraction.py b/youtube/yt_data_extract/watch_extraction.py
index a5ac0f5..72b6f2c 100644
--- a/youtube/yt_data_extract/watch_extraction.py
+++ b/youtube/yt_data_extract/watch_extraction.py
@@ -111,14 +111,10 @@ _formats = {
     '_rtmp': {'protocol': 'rtmp'},
 
     # av01 video only formats sometimes served with "unknown" codecs
-    '394': {'ext': 'mp4', 'height': 144, 'format_note': 'DASH video', 'vcodec': 'av01.0.00M.08'},
-    '395': {'ext': 'mp4', 'height': 240, 'format_note': 'DASH video', 'vcodec': 'av01.0.00M.08'},
-    '396': {'ext': 'mp4', 'height': 360, 'format_note': 'DASH video', 'vcodec': 'av01.0.01M.08'},
-    '397': {'ext': 'mp4', 'height': 480, 'format_note': 'DASH video', 'vcodec': 'av01.0.04M.08'},
-    '398': {'ext': 'mp4', 'height': 720, 'format_note': 'DASH video', 'vcodec': 'av01.0.05M.08'},
-    '399': {'ext': 'mp4', 'height': 1080, 'format_note': 'DASH video', 'vcodec': 'av01.0.08M.08'},
-    '400': {'ext': 'mp4', 'height': 1440, 'format_note': 'DASH video', 'vcodec': 'av01.0.12M.08'},
-    '401': {'ext': 'mp4', 'height': 2160, 'format_note': 'DASH video', 'vcodec': 'av01.0.12M.08'},
+    '394': {'vcodec': 'av01.0.05M.08'},
+    '395': {'vcodec': 'av01.0.05M.08'},
+    '396': {'vcodec': 'av01.0.05M.08'},
+    '397': {'vcodec': 'av01.0.05M.08'},
 }
 
 
@@ -137,29 +133,52 @@ def _extract_from_video_information_renderer(renderer_content):
     return info
 
 def _extract_likes_dislikes(renderer_content):
-    info = {
-        'like_count': None,
-    }
-    for button in renderer_content.get('buttons', ()):
-        button_renderer = button.get('slimMetadataToggleButtonRenderer', {})
-
+    def extract_button_count(toggle_button_renderer):
         # all the digits can be found in the accessibility data
-        count = extract_int(deep_get(
-                    button_renderer,
-                    'button', 'toggleButtonRenderer', 'defaultText',
-                    'accessibility', 'accessibilityData', 'label'))
+        count = extract_int(multi_deep_get(
+            toggle_button_renderer,
+            ['defaultText', 'accessibility', 'accessibilityData', 'label'],
+            ['accessibility', 'label'],
+            ['accessibilityData', 'accessibilityData', 'label'],
+        ))
 
         # this count doesn't have all the digits, it's like 53K for instance
         dumb_count = extract_int(extract_str(deep_get(
-            button_renderer, 'button', 'toggleButtonRenderer', 'defaultText')))
+            toggle_button_renderer, 'defaultText')))
 
         # The accessibility text will be "No likes" or "No dislikes" or
         # something like that, but dumb count will be 0
         if dumb_count == 0:
             count = 0
+        return count
 
-        if 'isLike' in button_renderer:
-            info['like_count'] = count
+    info = {
+        'like_count': None,
+        'dislike_count': None,
+    }
+    for button in renderer_content.get('buttons', ()):
+        if 'slimMetadataToggleButtonRenderer' in button:
+            button_renderer = button['slimMetadataToggleButtonRenderer']
+            count = extract_button_count(deep_get(button_renderer,
+                                                  'button',
+                                                  'toggleButtonRenderer'))
+            if 'isLike' in button_renderer:
+                info['like_count'] = count
+            elif 'isDislike' in button_renderer:
+                info['dislike_count'] = count
+        elif 'slimMetadataButtonRenderer' in button:
+            print(button)
+            button_renderer = button['slimMetadataButtonRenderer']
+            liberal_update(info, 'like_count', extract_button_count(deep_get(
+                button_renderer, 'button',
+                'segmentedLikeDislikeButtonRenderer',
+                'likeButton', 'toggleButtonRenderer'
+            )))
+            liberal_update(info, 'dislike_count',extract_button_count(deep_get(
+                button_renderer, 'button',
+                'segmentedLikeDislikeButtonRenderer',
+                'dislikeButton', 'toggleButtonRenderer'
+            )))
     return info
 
 def _extract_from_owner_renderer(renderer_content):
@@ -354,8 +373,10 @@ def _extract_watch_info_desktop(top_level):
     likes_dislikes = deep_get(video_info, 'sentimentBar', 'sentimentBarRenderer', 'tooltip', default='').split('/')
     if len(likes_dislikes) == 2:
         info['like_count'] = extract_int(likes_dislikes[0])
+        info['dislike_count'] = extract_int(likes_dislikes[1])
     else:
         info['like_count'] = None
+        info['dislike_count'] = None
 
     info['title'] = extract_str(video_info.get('title', None))
     info['author'] = extract_str(deep_get(video_info, 'owner', 'videoOwnerRenderer', 'title'))