touching up

Co-authored-by: bashonly <88596187+bashonly@users.noreply.github.com>

yt_dlp/extractor/drtv.py

@@ -3,6 +3,7 @@ import uuid
 
 from .common import InfoExtractor
 from ..utils import (
+    ExtractorError,
     int_or_none,
     mimetype2ext,
     parse_iso8601,
@@ -19,7 +20,7 @@ class DRTVIE(InfoExtractor):
     _VALID_URL = r'''(?x)
                     https?://
                         (?:
-                            (?:www\.)?dr\.dk/tv/se(?:/ondemand)?/(?:[^/]+/)*|
+                            (?:www\.)?dr\.dk/tv/se(?:/ondemand)?/(?:[^/?#]+/)*|
                             (?:www\.)?(?:dr\.dk|dr-massive\.com)/drtv/(?:se|episode|program)/
                         )
                         (?P<id>[\da-z_-]+)
@@ -125,45 +126,46 @@ class DRTVIE(InfoExtractor):
         'only_matching': True,
     }]
 
-    _TOKEN = None
-
     SUBTITLE_LANGS = {
         'DanishLanguageSubtitles': 'da',
         'ForeignLanguageSubtitles': 'da_foreign',
         'CombinedLanguageSubtitles': 'da_combined',
     }
 
+    _TOKEN = None
+
     def _real_initialize(self):
         if self._TOKEN:
             return
+
         token_response = self._download_json(
             'https://production.dr-massive.com/api/authorization/anonymous-sso', None,
             note='Downloading anonymous token', headers={
                 'content-type': 'application/json',
-            },
-            query={
+            }, query={
                 'device': 'web_browser',
                 'ff': 'idp,ldp,rpt',
                 'lang': 'da',
                 'supportFallbackToken': 'true',
-            },
-            data=json.dumps({
+            }, data=json.dumps({
                 'deviceId': str(uuid.uuid4()),
-                'scopes': [
-                    'Catalog',
-                ],
+                'scopes': ['Catalog'],
                 'optout': True,
             }).encode())
 
         self._TOKEN = traverse_obj(
             token_response, (lambda _, x: x['type'] == 'UserAccount', 'value', {str}), get_all=False)
+        if not self._TOKEN:
+            raise ExtractorError('Unable to get anonymous token')
 
     def _real_extract(self, url):
         url_slug = self._match_id(url)
         webpage = self._download_webpage(url, url_slug)
 
-        json_data = self._search_json(r'window\.__data\s*=\s*', webpage, 'data', url_slug, fatal=False) or {}
-        item = traverse_obj(json_data, ('cache', 'page', ..., (None, ('entries', 0)), 'item'), get_all=False)
+        json_data = self._search_json(
+            r'window\.__data\s*=', webpage, 'data', url_slug, fatal=False) or {}
+        item = traverse_obj(
+            json_data, ('cache', 'page', ..., (None, ('entries', 0)), 'item', {dict}), get_all=False)
         if item:
             item_id = item.get('id')
         else:
@@ -180,6 +182,7 @@ class DRTVIE(InfoExtractor):
                     'segments': 'drtv,optedout',
                     'sub': 'Anonymous',
                 })
+
         video_id = try_call(lambda: item['customId'].rsplit(':', 1)[-1]) or item_id
         stream_data = self._download_json(
             f'https://production.dr-massive.com/api/account/items/{item_id}/videos', video_id,
@@ -190,16 +193,13 @@ class DRTVIE(InfoExtractor):
                 'lang': 'da',
                 'resolution': 'HD-1080',
                 'sub': 'Anonymous',
-            },
-            headers={
-                'authorization': f'Bearer {self._TOKEN}',
-            })
+            }, headers={'authorization': f'Bearer {self._TOKEN}'})
 
         formats = []
         subtitles = {}
-        for fmt in traverse_obj(stream_data, (lambda _, x: x['url'], {dict})):
-            format_id = fmt.get('format', 'na')
-            access_service = fmt.get('accessService')
+        for stream in traverse_obj(stream_data, (lambda _, x: x['url'])):
+            format_id = stream.get('format', 'na')
+            access_service = stream.get('accessService')
             preference = None
             subtitle_suffix = ''
             if access_service in ('SpokenSubtitles', 'SignLanguage', 'VisuallyInterpreted'):
@@ -209,10 +209,10 @@ class DRTVIE(InfoExtractor):
             elif access_service == 'StandardVideo':
                 preference = 1
             fmts, subs = self._extract_m3u8_formats_and_subtitles(
-                fmt.get('url'), video_id, preference=preference, m3u8_id=format_id, fatal=False)
+                stream.get('url'), video_id, preference=preference, m3u8_id=format_id, fatal=False)
             formats.extend(fmts)
 
-            api_subtitles = traverse_obj(fmt, ('subtitles', lambda _, v: url_or_none(v['link']), {dict}))
+            api_subtitles = traverse_obj(stream, ('subtitles', lambda _, v: url_or_none(v['link']), {dict}))
             if not api_subtitles:
                 self._merge_subtitles(subs, target=subtitles)
 
@@ -223,6 +223,9 @@ class DRTVIE(InfoExtractor):
                     'ext': mimetype2ext(sub_track.get('format')) or 'vtt'
                 })
 
+        if not formats and traverse_obj(item, ('season', 'customFields', 'IsGeoRestricted')):
+            self.raise_geo_restricted(countries=self._GEO_COUNTRIES)
+
         return {
             'id': video_id,
             'formats': formats,